Merge pull request #91 from BillFarber/task/extendExamples

Adds a Vector/BM25 example.

Author: BillFarber

examples/langchain/README.md

@@ -104,7 +104,9 @@ loaded via `load_data.py`.
 ### MarkLogic 12EA Setup
 
 To try out this functionality out, you will need acces to an instance of MarkLogic 12
-(currently internal or Early Access only). You may use docker 
+(currently internal or Early Access only).
+<TODO>Add info to get ML12</TODO>
+You may use docker 
 [docker-compose](https://docs.docker.com/compose/) to instantiate a new MarkLogic
 instance with port 8003 available (you can use your own MarkLogic instance too, just be
 sure that port 8003 is available):
@@ -147,3 +149,18 @@ into different collections.
 ```
 python load_data_with_embeddings.py
 ```
+
+### Running the Vector Query
+
+You are now ready to test the example vector retriever. Run the following to ask a
+question with the results augmented via the `marklogic_vector_query_retriever.py` module
+in this project:
+
+    python ask_vector_query.py "What is task decomposition?" posts_with_embeddings
+
+This retriever searches MarkLogic for candidate documents, and defaults to
+using the new score-bm25 scoring method in MarkLogic 12EA. If preferred, you can adjust
+this to one of the other scoring methods. After retrieving candidate documents based on
+the CTS search, the retriever uses the new vector functionality to sort the documents
+based on cosine similarity to the user question, and then returns the top N documents
+for the retriever to package up.

examples/langchain/ask_vector_query.py

@@ -0,0 +1,53 @@
+# Based on example at
+# https://python.langchain.com/docs/use_cases/question_answering/quickstart .
+
+import os
+import sys
+from dotenv import load_dotenv
+from langchain import hub
+from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
+from langchain.schema import StrOutputParser
+from langchain.schema.runnable import RunnablePassthrough
+from marklogic import Client
+from marklogic_vector_query_retriever import (
+    MarkLogicVectorQueryRetriever,
+)
+
+
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+
+load_dotenv()
+embeddings = AzureOpenAIEmbeddings(
+    azure_deployment=os.environ["AZURE_EMBEDDING_DEPLOYMENT_NAME"]
+)
+retriever = MarkLogicVectorQueryRetriever.create(
+    Client("http://localhost:8003", digest=("langchain-user", "password")),
+    embedding_generator=embeddings,
+)
+retriever.collections = [sys.argv[2]]
+retriever.max_results = int(sys.argv[3]) if len(sys.argv) > 3 else 10
+if len(sys.argv) > 4:
+    retriever.query_type = sys.argv[4]
+
+question = sys.argv[1]
+
+prompt = hub.pull("rlm/rag-prompt")
+# Note that the Azure OpenAI API key, the Azure OpenAI Endpoint, and the OpenAI API
+# Version, are all read from the environment automatically.
+llm = AzureChatOpenAI(
+    model_name=os.getenv("AZURE_LLM_DEPLOYMENT_NAME"),
+    azure_deployment=os.getenv("AZURE_LLM_DEPLOYMENT_NAME"),
+    temperature=0,
+    max_tokens=None,
+    timeout=None,
+)
+
+rag_chain = (
+    {"context": retriever | format_docs, "question": RunnablePassthrough()}
+    | prompt
+    | llm
+    | StrOutputParser()
+)
+print(rag_chain.invoke(question))

examples/langchain/load_data.py

@@ -17,34 +17,42 @@
 )
 docs = loader.load()
 
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=1000, chunk_overlap=100
+)
 splits = text_splitter.split_documents(docs)
 
 client = Client("http://localhost:8003", digest=("langchain-user", "password"))
 
-marklogic_docs = [
-    DefaultMetadata(collections="posts")
-]
+marklogic_docs = [DefaultMetadata(collections="posts")]
 for split in splits:
-    doc = Document(None, split.page_content, extension=".txt", directory="/post/")
+    doc = Document(
+        None, split.page_content, extension=".txt", directory="/post/"
+    )
     marklogic_docs.append(doc)
 
 client.documents.write(marklogic_docs)
-print(f"Number of documents written to collection 'posts': {len(marklogic_docs)-1}")
+print(
+    f"Number of documents written to collection 'posts': {len(marklogic_docs)-1}"
+)
 
 loader = WebBaseLoader(
-    web_paths=("https://raw.githubusercontent.com/langchain-ai/langchain/master/docs/docs/modules/state_of_the_union.txt",)
+    web_paths=(["https://www.whitehouse.gov/state-of-the-union-2022/"])
 )
 docs = loader.load()
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=1000, chunk_overlap=100
+)
 splits = text_splitter.split_documents(docs)
 
-marklogic_docs = [
-    DefaultMetadata(collections="sotu")
-]
+marklogic_docs = [DefaultMetadata(collections="sotu")]
 for split in splits:
-    doc = Document(None, split.page_content, extension=".txt", directory="/sotu/")
+    doc = Document(
+        None, split.page_content, extension=".txt", directory="/sotu/"
+    )
     marklogic_docs.append(doc)
 
 client.documents.write(marklogic_docs)
-print(f"Number of documents written to collection 'sotu': {len(marklogic_docs)-1}")
+print(
+    f"Number of documents written to collection 'sotu': {len(marklogic_docs)-1}"
+)

examples/langchain/load_data_with_embeddings.py

@@ -50,9 +50,7 @@
 )
 
 loader = WebBaseLoader(
-    web_paths=(
-        "https://raw.githubusercontent.com/langchain-ai/langchain/master/docs/docs/modules/state_of_the_union.txt",
-    )
+    web_paths=(["https://www.whitehouse.gov/state-of-the-union-2022/"])
 )
 docs = loader.load()
 text_splitter = RecursiveCharacterTextSplitter(

examples/langchain/marklogic_vector_query_retriever.py

@@ -0,0 +1,92 @@
+from typing import List
+from langchain_core.documents import Document
+from langchain_core.retrievers import BaseRetriever
+from langchain_openai import AzureOpenAIEmbeddings
+from marklogic import Client
+
+
+class MarkLogicVectorQueryRetriever(BaseRetriever):
+
+    client: Client
+    embedding_generator: AzureOpenAIEmbeddings
+    max_results: int = 10
+    collections: List[str] = []
+    tde_schema: str
+    tde_view: str
+    scoring_method: str
+
+    @classmethod
+    def create(
+        cls,
+        client: Client,
+        embedding_generator: AzureOpenAIEmbeddings,
+        tde_schema: str = None,
+        tde_view: str = None,
+        scoring_method: str = "score-bm25",
+    ):
+        return cls(
+            client=client,
+            embedding_generator=embedding_generator,
+            tde_schema=tde_schema or "demo",
+            tde_view=tde_view or "posts",
+            scoring_method=scoring_method,
+        )
+
+    def _build_javascript_query_query(self, query, query_embedding):
+        # Returning first self.max_results documents based on token limitations
+        #
+        # If limits are hit, consider different models:
+        # gpt-35-turbo (0125): 16,385/4,096
+        # gpt-35-turbo (1106): 16,385/4,096
+        # gpt-35-turbo-16k (0613):
+
+        # This JavaScript consists of two queries.
+        # The first is a standard cts search, searching for words that match those used
+        # in the chat question.
+        # The second query is an Optic query that uses the top documents from the CTS
+        # query to do a vector search to re-order the results.
+
+        search_words = []
+        for word in query.split():
+            search_words.append(word.lower().replace("?", ""))
+        return """
+            const op = require('/MarkLogic/optic');
+            const ovec = require('/MarkLogic/optic/optic-vec.xqy');
+            const result =
+                fn.subsequence(cts.search(cts.andQuery([
+                    cts.wordQuery({}),
+                    cts.collectionQuery({})
+                ]),["{}"]), 1, {});
+            let uris = [];
+            for(const doc of result){{
+                uris.push(xdmp.nodeUri(doc))
+            }}
+            const qv = vec.vector({})
+
+            const rows = op.fromView('{}','{}','')
+            .where(op.in(op.col('uri'), uris))
+            .bind(op.as('summaryCosineSim', op.vec.cosineSimilarity(op.vec.vector(op.col('embedding')),qv)))
+            .orderBy(op.desc(op.col('summaryCosineSim')))
+            .result();
+            rows;
+        """.format(
+            search_words,
+            self.collections,
+            self.scoring_method,
+            self.max_results,
+            query_embedding,
+            self.tde_schema,
+            self.tde_view,
+        )
+
+    def _get_relevant_documents(self, query: str) -> List[Document]:
+        print(f"Searching with query: {query}")
+
+        query_embedding = self.embedding_generator.embed_query(query)
+        javascript_query = self._build_javascript_query_query(
+            query, query_embedding
+        )
+        results = self.client.eval(javascript=javascript_query)
+
+        print(f"Count of matching MarkLogic documents: {len(results)}")
+        return map(lambda doc: Document(page_content=doc["text"]), results)

examples/langchain/src/main/ml-config/security/roles/langchain-eval-role.json

@@ -0,0 +1,20 @@
+{
+  "role-name": "langchain-eval-role",
+  "privilege": [
+    {
+      "privilege-name": "xdmp:eval",
+      "action": "http://marklogic.com/xdmp/privileges/xdmp-eval",
+      "kind": "execute"
+    },
+    {
+      "privilege-name": "xdmp:eval-in",
+      "action": "http://marklogic.com/xdmp/privileges/xdmp-eval-in",
+      "kind": "execute"
+    },
+    {
+      "privilege-name": "xdbc:eval",
+      "action": "http://marklogic.com/xdmp/privileges/xdbc-eval",
+      "kind": "execute"
+    }
+  ]
+}

examples/langchain/src/main/ml-config/security/users/langchain-user.json

@@ -4,6 +4,7 @@
     "role": [
         "rest-reader",
         "rest-writer",
-        "qconsole-user"
+        "qconsole-user",
+        "langchain-eval-role"
     ]
 }