Added doc embedding support (#41)

* document embedding writer & query
* Added test query for doc embeddings
* Bump version
* Added doc rag prompt
* Document RAG service

trustgraph/model/prompt/generic/prompts.py

@@ -79,3 +79,18 @@ Use only the provided knowledge statements to respond to the following:
 """
 
     return prompt
+
+def to_document_query(query, documents):
+
+    documents = "\n\n".join(documents)
+
+    prompt=f"""Study the following context. Use only the information provided in the context in your response. Do not speculate if the answer is not found in the provided set of knowledge statements.
+
+Here is the context:
+{documents}
+
+Use only the provided knowledge statements to respond to the following:
+{query}
+"""
+
+    return prompt

trustgraph/model/prompt/generic/service.py

@@ -1,4 +1,3 @@
-
 """
 Language service abstracts prompt engineering from LLM.
 """
@@ -14,7 +13,8 @@ from .... schema import prompt_request_queue, prompt_response_queue
 from .... base import ConsumerProducer
 from .... clients.llm_client import LlmClient
 
-from . prompts import to_definitions, to_relationships, to_kg_query
+from . prompts import to_definitions, to_relationships
+from . prompts import to_kg_query, to_document_query
 
 module = ".".join(__name__.split(".")[1:-1])
 
@@ -82,6 +82,11 @@ class Processor(ConsumerProducer):
             self.handle_kg_prompt(id, v)
             return
 
+        elif kind == "document-prompt":
+
+            self.handle_document_prompt(id, v)
+            return
+
         else:
 
             print("Invalid kind.", flush=True)
@@ -252,6 +257,43 @@ class Processor(ConsumerProducer):
 
             self.producer.send(r, properties={"id": id})
         
+    def handle_document_prompt(self, id, v):
+
+        try:
+
+            prompt = to_document_query(v.query, v.documents)
+
+            print("prompt")
+            print(prompt)
+
+            print("Call LLM...")
+
+            ans = self.llm.request(prompt)
+
+            print(ans)
+
+            print("Send response...", flush=True)
+            r = PromptResponse(answer=ans, error=None)
+            self.producer.send(r, properties={"id": id})
+
+            print("Done.", flush=True)
+
+        except Exception as e:
+
+            print(f"Exception: {e}")
+
+            print("Send error response...", flush=True)
+
+            r = PromptResponse(
+                error=Error(
+                    type = "llm-error",
+                    message = str(e),
+                ),
+                response=None,
+            )
+
+            self.producer.send(r, properties={"id": id})
+        
     @staticmethod
     def add_args(parser):
 

trustgraph/model/prompt/template/prompts.py

@@ -17,3 +17,7 @@ def to_kg_query(template, query, kg):
     cypher =  get_cypher(kg)
     return template.format(query=query, graph=cypher)
 
+def to_document_query(template, query, docs):
+    docs = "\n\n".join(docs)
+    return template.format(query=query, documents=docs)
+

trustgraph/model/prompt/template/service.py

@@ -14,7 +14,8 @@ from .... schema import prompt_request_queue, prompt_response_queue
 from .... base import ConsumerProducer
 from .... clients.llm_client import LlmClient
 
-from . prompts import to_definitions, to_relationships, to_kg_query
+from . prompts import to_definitions, to_relationships
+from . prompts import to_kg_query, to_document_query
 
 module = ".".join(__name__.split(".")[1:-1])
 
@@ -38,6 +39,7 @@ class Processor(ConsumerProducer):
         definition_template = params.get("definition_template")
         relationship_template = params.get("relationship_template")
         knowledge_query_template = params.get("knowledge_query_template")
+        document_query_template = params.get("document_query_template")
 
         super(Processor, self).__init__(
             **params | {
@@ -48,9 +50,6 @@ class Processor(ConsumerProducer):
                 "output_schema": PromptResponse,
                 "text_completion_request_queue": tc_request_queue,
                 "text_completion_response_queue": tc_response_queue,
-                "definition_template": definition_template,
-                "relationship_template": relationship_template,
-                "knowledge_query_template": knowledge_query_template,
             }
         )
 
@@ -64,6 +63,7 @@ class Processor(ConsumerProducer):
         self.definition_template = definition_template
         self.relationship_template = relationship_template
         self.knowledge_query_template = knowledge_query_template
+        self.document_query_template = document_query_template
 
     def handle(self, msg):
 
@@ -92,6 +92,11 @@ class Processor(ConsumerProducer):
             self.handle_kg_prompt(id, v)
             return
 
+        elif kind == "document-prompt":
+
+            self.handle_document_prompt(id, v)
+            return
+
         else:
 
             print("Invalid kind.", flush=True)
@@ -120,6 +125,11 @@ class Processor(ConsumerProducer):
                     e = defn["entity"]
                     d = defn["definition"]
 
+                    if e == "": continue
+                    if e is None: continue
+                    if d == "": continue
+                    if d is None: continue
+
                     output.append(
                         Definition(
                             name=e, definition=d
@@ -171,12 +181,30 @@ class Processor(ConsumerProducer):
             for defn in defs:
 
                 try:
+
+                    s = defn["subject"]
+                    p = defn["predicate"]
+                    o = defn["object"]
+                    o_entity = defn["object-entity"]
+
+                    if s == "": continue
+                    if s is None: continue
+
+                    if p == "": continue
+                    if p is None: continue
+
+                    if o == "": continue
+                    if o is None: continue
+
+                    if o_entity == "" or o_entity is None:
+                        o_entity = False
+
                     output.append(
                         Relationship(
-                            s = defn["subject"],
-                            p = defn["predicate"],
-                            o = defn["object"],
-                            o_entity = defn["object-entity"],
+                            s = s,
+                            p = p,
+                            o = o,
+                            o_entity = o_entity,
                         )
                     )
 
@@ -239,6 +267,42 @@ class Processor(ConsumerProducer):
 
             self.producer.send(r, properties={"id": id})
         
+    def handle_document_prompt(self, id, v):
+
+        try:
+
+            prompt = to_document_query(
+                self.document_query_template, v.query, v.documents
+            )
+
+            print(prompt)
+
+            ans = self.llm.request(prompt)
+
+            print(ans)
+
+            print("Send response...", flush=True)
+            r = PromptResponse(answer=ans, error=None)
+            self.producer.send(r, properties={"id": id})
+
+            print("Done.", flush=True)
+
+        except Exception as e:
+
+            print(f"Exception: {e}")
+
+            print("Send error response...", flush=True)
+
+            r = PromptResponse(
+                error=Error(
+                    type = "llm-error",
+                    message = str(e),
+                ),
+                response=None,
+            )
+
+            self.producer.send(r, properties={"id": id})
+        
     @staticmethod
     def add_args(parser):
 
@@ -277,6 +341,12 @@ class Processor(ConsumerProducer):
             help=f'Knowledge query template',
         )
 
+        parser.add_argument(
+            '--document-query-template',
+            required=True,
+            help=f'Document query template',
+        )
+
 def run():
 
     Processor.start(module, __doc__)