Doc RAG working

trustgraph-base/trustgraph/base/__init__.py

@@ -25,4 +25,5 @@ from . graph_embeddings_query_service import GraphEmbeddingsQueryService
 from . document_embeddings_query_service import DocumentEmbeddingsQueryService
 from . graph_embeddings_client import GraphEmbeddingsClientSpec
 from . triples_client import TriplesClientSpec
+from . document_embeddings_client import DocumentEmbeddingsClientSpec
 

trustgraph-base/trustgraph/base/document_embeddings_client.py

@@ -0,0 +1,38 @@
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import DocumentEmbeddingsRequest, DocumentEmbeddingsResponse
+from .. knowledge import Uri, Literal
+
+class DocumentEmbeddingsClient(RequestResponse):
+    async def query(self, vectors, limit=20, user="trustgraph",
+                    collection="default", timeout=30):
+
+        resp = await self.request(
+            DocumentEmbeddingsRequest(
+                vectors = vectors,
+                limit = limit,
+                user = user,
+                collection = collection
+            ),
+            timeout=timeout
+        )
+
+        print(resp, flush=True)
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        return resp.documents
+
+class DocumentEmbeddingsClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(DocumentEmbeddingsClientSpec, self).__init__(
+            request_name = request_name,
+            request_schema = DocumentEmbeddingsRequest,
+            response_name = response_name,
+            response_schema = DocumentEmbeddingsResponse,
+            impl = DocumentEmbeddingsClient,
+        )
+

trustgraph-base/trustgraph/base/document_embeddings_query_service.py

@@ -49,10 +49,10 @@ class DocumentEmbeddingsQueryService(FlowProcessor):
 
             print(f"Handling input {id}...", flush=True)
 
-            entities = self.query_document_embeddings(request)
+            docs = await self.query_document_embeddings(request)
 
             print("Send response...", flush=True)
-            r = DocumentEmbeddingsResponse(entities=entities, error=None)
+            r = DocumentEmbeddingsResponse(documents=docs, error=None)
             await flow("response").send(r, properties={"id": id})
 
             print("Done.", flush=True)

trustgraph-base/trustgraph/base/prompt_client.py

@@ -53,6 +53,16 @@ class PromptClient(RequestResponse):
             timeout = timeout,
         )
 
+    async def document_prompt(self, query, documents, timeout=600):
+        return await self.prompt(
+            id = "document-prompt",
+            variables = {
+                "query": query,
+                "documents": documents,
+            },
+            timeout = timeout,
+        )
+
 class PromptClientSpec(RequestResponseSpec):
     def __init__(
             self, request_name, response_name,

trustgraph-base/trustgraph/schema/documents.py

@@ -11,8 +11,6 @@ class Document(Record):
     metadata = Metadata()
     data = Bytes()
 
-document_ingest_queue = topic('document-load')
-
 ############################################################################
 
 # Text documents / text from PDF
@@ -21,8 +19,6 @@ class TextDocument(Record):
     metadata = Metadata()
     text = Bytes()
 
-text_ingest_queue = topic('text-document-load')
-
 ############################################################################
 
 # Chunks of text
@@ -31,8 +27,6 @@ class Chunk(Record):
     metadata = Metadata()
     chunk = Bytes()
 
-chunk_ingest_queue = topic('chunk-load')
-
 ############################################################################
 
 # Document embeddings are embeddings associated with a chunk
@@ -46,8 +40,6 @@ class DocumentEmbeddings(Record):
     metadata = Metadata()
     chunks = Array(ChunkEmbeddings())
 
-document_embeddings_store_queue = topic('document-embeddings-store')
-
 ############################################################################
 
 # Doc embeddings query
@@ -62,10 +54,3 @@ class DocumentEmbeddingsResponse(Record):
     error = Error()
     documents = Array(Bytes())
 
-document_embeddings_request_queue = topic(
-    'doc-embeddings', kind='non-persistent', namespace='request'
-)
-document_embeddings_response_queue = topic(
-    'doc-embeddings', kind='non-persistent', namespace='response',
-)
-