Refactor Corpus instantiation flow, as well as API data structure (#9)

Author: maxyu1115

format.sh

@@ -1,4 +1,4 @@
-directories="memas tests integration-tests memas_client memas_sdk benchmarking"
+directories="memas tests integration-tests"
 for dir in $directories; do
     find $dir -type f -name "*.py" -exec autopep8 --max-line-length 120 -i {} \;
 done

integration-tests/corpus/test_basic_corpus.py

@@ -1,24 +1,22 @@
 import numpy as np
 import uuid
 import time
-from memas.interface.storage_driver import DocumentEntity
-from memas.storage_driver.corpus_vector_store import CorpusVectorStore
-from memas.storage_driver import corpus_vector_store, corpus_doc_store, corpus_doc_metadata
+from memas.context_manager import ctx
 from memas.corpus import basic_corpus
 from memas.interface.corpus import Citation
 
 corpus_name = "test corpus1"
-test_corpus = basic_corpus.BasicCorpus(uuid.uuid4(), corpus_name)
-
 
 def test_save_then_search_one_corpus(es_client):
+    test_corpus = basic_corpus.BasicCorpus(uuid.uuid4(), corpus_name, ctx.corpus_metadata, ctx.corpus_doc, ctx.corpus_vec)
+
     text1 = "The sun is high. California sunshine is great. "
     text2 = "I picked up my phone and then dropped it again. I cant seem to get a good grip on things these days. It persists into my everyday tasks"
     text3 = "The weather is great today, but I worry that tomorrow it won't be. My umbrella is in the repair shop."
 
-    assert test_corpus.store_and_index(text1, "doc1", Citation("www.docsource1", "SSSdoc1", ""))
-    assert test_corpus.store_and_index(text2, "doc2", Citation("were.docsource2", "SSSdoc2", ""))
-    assert test_corpus.store_and_index(text3, "doc3", Citation("docsource3.ai", "SSSdoc3", ""))
+    assert test_corpus.store_and_index(text1, Citation("www.docsource1", "SSSdoc1", "", "doc1"))
+    assert test_corpus.store_and_index(text2, Citation("were.docsource2", "SSSdoc2", "", "doc2"))
+    assert test_corpus.store_and_index(text3, Citation("docsource3.ai", "SSSdoc3", "", "doc3"))
 
     time.sleep(1)
     output = test_corpus.search("It is sunny")

integration-tests/storage_driver/test_corpus_doc_metadata.py

@@ -14,6 +14,6 @@ def test_insert_and_get():
     corpus_id = uuid.uuid4()
     document_id = uuid.uuid4()
 
-    citation = Citation("google.com", "test google", "just a simple test")
-    metadata.insert_document_metadata(corpus_id, document_id, 1, "test", citation)
+    citation = Citation("google.com", "test google", "just a simple test", "test")
+    metadata.insert_document_metadata(corpus_id, document_id, 1, citation)
     assert metadata.get_document_citation(corpus_id, document_id) == citation

memas/app.py

@@ -1,8 +1,6 @@
 import yaml
 from flask import Flask
 from memas.context_manager import ContextManager
-from memas.interface.corpus import CorpusType
-from memas.corpus.basic_corpus import BasicCorpusFactory
 
 
 def create_app(config_filename, *, first_init=False):
@@ -16,10 +14,6 @@ def create_app(config_filename, *, first_init=False):
 
     app.ctx.init()
 
-    # TODO : Need a better place to put this
-    app.ctx.corpus_provider.setCorpusFactory(CorpusType.KNOWLEDGE, BasicCorpusFactory())
-    app.ctx.corpus_provider.setCorpusFactory(CorpusType.CONVERSATION, BasicCorpusFactory())
-
     from memas.dataplane import dataplane
     from memas.controlplane import controlplane
     app.register_blueprint(dataplane)

memas/context_manager.py

@@ -71,7 +71,7 @@ def __init__(self, app_config: Config):
         self.es: Elasticsearch
 
         # Corpus provider
-        self.corpus_provider: CorpusProvider = CorpusProvider()
+        self.corpus_provider: CorpusProvider
 
     def setup_cassandra_keyspace(self):
         """Setup the cassandra keyspace. We only want to run the very first server launch. 
@@ -121,6 +121,8 @@ def init_datastores(self) -> None:
         self.corpus_vec.init()
         self.corpus_doc.init()
 
+        self.corpus_provider = CorpusProvider(self.corpus_metadata, self.corpus_doc, self.corpus_vec)
+
     def init(self) -> None:
         self.init_clients()
         self.init_datastores()

memas/corpus/basic_corpus.py

@@ -1,12 +1,10 @@
 # from search_redirect import SearchSettings
 import logging
 import uuid
-from functools import reduce
 from memas.interface.corpus import Corpus, CorpusFactory
 from memas.interface.corpus import Citation
-from memas.interface.storage_driver import DocumentEntity
+from memas.interface.storage_driver import CorpusDocumentMetadataStore, CorpusDocumentStore, CorpusVectorStore, DocumentEntity
 from memas.interface.exceptions import SentenceLengthOverflowException
-from memas.context_manager import ctx
 from memas.text_parsing.text_parsers import segment_document
 from memas.corpus.corpus_searching import normalize_and_combine
 
@@ -17,27 +15,30 @@
 
 class BasicCorpus(Corpus):
 
-    def __init__(self, corpus_id: uuid.UUID, corpus_name: str):
+    def __init__(self, corpus_id: uuid.UUID, corpus_name: str, metadata_store: CorpusDocumentMetadataStore, doc_store: CorpusDocumentStore, vec_store: CorpusVectorStore):
         super().__init__(corpus_id, corpus_name)
+        self.metadata_store: CorpusDocumentMetadataStore = metadata_store
+        self.doc_store: CorpusDocumentStore = doc_store
+        self.vec_store: CorpusVectorStore = vec_store
 
     """
     The function stores a document in the elastic search DB, vecDB, and doc MetaData.
     Returns True on Success, False on Failure
     """
 
-    def store_and_index(self, document: str, document_name: str, citation: Citation) -> bool:
+    def store_and_index(self, document: str, citation: Citation) -> bool:
         _log.debug(f"Corpus storing and indexing [corpus_id={self.corpus_id}]")
 
         doc_id = uuid.uuid4()
-        doc_entity = DocumentEntity(self.corpus_id, doc_id, document_name, document)
+        doc_entity = DocumentEntity(self.corpus_id, doc_id, citation.document_name, document)
 
         document_chunks = segment_document(document, MAX_SEGMENT_LENGTH)
 
         # TODO : Need to investigate how to undo when failures on partial insert
-        meta_save = ctx.corpus_metadata.insert_document_metadata(
-            self.corpus_id, doc_id, len(document_chunks), document_name, citation)
+        meta_save = self.metadata_store.insert_document_metadata(
+            self.corpus_id, doc_id, len(document_chunks), citation)
 
-        vec_save = ctx.corpus_vec.save_documents([doc_entity])
+        vec_save = self.vec_store.save_documents([doc_entity])
 
         # Divide longer documents for document store
         chunk_num = 0
@@ -46,11 +47,11 @@ def store_and_index(self, document: str, document_name: str, citation: Citation)
             # Create the new IDs for the document chunk combo
             chunk_id = doc_id.hex + '{:032b}'.format(chunk_num)
             chunk_num = chunk_num + 1
-            doc_chunk_entity = DocumentEntity(self.corpus_id, doc_id, document_name, chunk)
+            doc_chunk_entity = DocumentEntity(self.corpus_id, doc_id, citation.document_name, chunk)
             chunk_id_entity_pairs.append((chunk_id, doc_chunk_entity))
 
         # Insert all chunks of document at once
-        doc_save = ctx.corpus_doc.save_documents(id_doc_pairs=chunk_id_entity_pairs)
+        doc_save = self.doc_store.save_documents(id_doc_pairs=chunk_id_entity_pairs)
 
         return meta_save and vec_save and doc_save
 
@@ -67,24 +68,24 @@ def search(self, clue: str) -> list[tuple[float, str, Citation]]:
         vector_search_count: int = 10
 
         doc_store_results: list[tuple[float, str, Citation]] = []
-        temp_res = ctx.corpus_doc.search_corpora([self.corpus_id], clue)
+        temp_res = self.doc_store.search_corpora([self.corpus_id], clue)
         # Search the document store
         for score, doc_entity in temp_res:
             document_text = doc_entity.document
-            citation = ctx.corpus_metadata.get_document_citation(self.corpus_id, doc_entity.document_id)
+            citation = self.metadata_store.get_document_citation(self.corpus_id, doc_entity.document_id)
             doc_store_results.append([score, document_text, citation])
 
         # Search for the vectors
         vec_store_results: list[tuple[float, str, Citation]] = []
-        temp_res2 = ctx.corpus_vec.search_corpora([self.corpus_id], clue)
+        temp_res2 = self.vec_store.search_corpora([self.corpus_id], clue)
         for score, doc_entity, start_index, end_index in temp_res2:
 
             # Verify that the text recovered from the vectors fits the maximum sentence criteria
             if end_index - start_index != len(doc_entity.document):
                 _log.error("Index not aligned with actual document", exc_info=True)
                 raise SentenceLengthOverflowException(end_index - start_index)
 
-            citation = ctx.corpus_metadata.get_document_citation(self.corpus_id, doc_entity.document_id)
+            citation = self.metadata_store.get_document_citation(self.corpus_id, doc_entity.document_id)
             vec_store_results.append([score, doc_entity.document, citation])
 
         # If any of the searches returned no results combine and return
@@ -100,11 +101,14 @@ def search(self, clue: str) -> list[tuple[float, str, Citation]]:
 
         return results
 
-    def generate_search_instructions(self, clue: str) -> any:
-        pass
-
 
 class BasicCorpusFactory(CorpusFactory):
+    def __init__(self, metadata_store: CorpusDocumentMetadataStore, doc_store: CorpusDocumentStore, vec_store: CorpusVectorStore) -> None:
+        super().__init__()
+        self.metadata_store: CorpusDocumentMetadataStore = metadata_store
+        self.doc_store: CorpusDocumentStore = doc_store
+        self.vec_store: CorpusVectorStore = vec_store
+    
     def produce(self, corpus_id: uuid.UUID):
         # TODO: Maybe change the Corpus Name Parameter
-        return BasicCorpus(corpus_id, "BasicCorpus")
+        return BasicCorpus(corpus_id, "BasicCorpus", self.metadata_store, self.doc_store, self.vec_store)

memas/corpus/corpus_provider.py

@@ -1,27 +1,29 @@
 import logging
 from uuid import UUID
+from memas.corpus.basic_corpus import BasicCorpusFactory
 from memas.interface.corpus import Corpus, CorpusFactory, CorpusType
+from memas.interface.storage_driver import CorpusDocumentMetadataStore, CorpusDocumentStore, CorpusVectorStore
 
 
 _log = logging.getLogger(__name__)
 
 
 class CorpusProvider:
-    def __init__(self) -> None:
+    def __init__(self, metadata_store: CorpusDocumentMetadataStore, doc_store: CorpusDocumentStore, vec_store: CorpusVectorStore) -> None:
         self.factory_dict: dict[CorpusType, CorpusFactory] = dict()
+        
+        basic_corpus_factory = BasicCorpusFactory(metadata_store, doc_store, vec_store)
+        self.factory_dict[CorpusType.CONVERSATION] = basic_corpus_factory
+        self.factory_dict[CorpusType.KNOWLEDGE] = basic_corpus_factory
 
-    def setCorpusFactory(self, corpus_type: CorpusType, corpus_factory: CorpusFactory):
-        self.factory_dict[corpus_type] = corpus_factory
-
-    # TODO : Fix the last parameter that was just removed - what is that supposed to be for? namespace_id
-
-    def get_corpus(self, corpus_id: UUID, *, corpus_type: CorpusType) -> Corpus:
+    def get_corpus(self, corpus_id: UUID, *, corpus_type: CorpusType, namespace_id: UUID=None) -> Corpus:
         """Gets the Corpus class based on the corpus_id
 
         Args:
             corpus_id (UUID): corpus_id
             corpus_type (CorpusType): type of the corpus, this is necessary unless a namespace_id is provided
-            namespace_id (UUID): namespace_id of the corpus. This is necessary when
+            namespace_id (UUID): namespace_id of the corpus, this is necessary when a corpus type is not provided, 
+                since it's needed to find the corpus type.
 
         Returns:
             Corpus: _description_

memas/corpus/corpus_searching.py

@@ -5,7 +5,6 @@
 from memas.interface.corpus import Citation
 from memas.interface.storage_driver import DocumentEntity
 from memas.interface.exceptions import SentenceLengthOverflowException
-from memas.context_manager import ctx
 
 
 def corpora_search(corpus_ids: list[UUID], clue: str) -> list[tuple[float, str, Citation]]:

memas/dataplane.py

@@ -2,9 +2,7 @@
 from flask import Blueprint, current_app, request
 from memas.context_manager import ctx
 from memas.interface.corpus import Citation, Corpus, CorpusType
-from memas.storage_driver.memas_metadata import split_corpus_pathname
-from memas.corpus.basic_corpus import BasicCorpusFactory
-from memas.corpus.corpus_searching import corpora_search
+
 
 dataplane = Blueprint("dp", __name__, url_prefix="/dp")
 
@@ -38,20 +36,21 @@ def recall():
 def memorize():
     corpus_pathname: str = request.json["corpus_pathname"]
     document: str = request.json["document"]
-    document_name: str = request.json.get("document_name", "")
+    raw_citation: str = request.json["citation"]
+
+    document_name = raw_citation.get("document_name", "")
 
     current_app.logger.info(f"Memorizing [corpus_pathname=\"{corpus_pathname}\"] [document_name=\"{document_name}\"]")
 
-    # TODO : need to be able to fetch the corpus name for citation purposes
-    corpus_name = split_corpus_pathname(corpus_pathname)[1]
-    raw_citation: str = request.json["citation"]
-    citation = Citation(raw_citation["source_uri"], raw_citation["source_name"],
-                        raw_citation["description"])
+    citation = Citation(source_uri=raw_citation.get("source_uri", ""),
+                        source_name=raw_citation.get("source_name", ""),
+                        description=raw_citation.get("description", ""),
+                        document_name=document_name)
 
     corpus_info = ctx.memas_metadata.get_corpus_info(corpus_pathname)
 
     corpus: Corpus = ctx.corpus_provider.get_corpus(corpus_info.corpus_id, corpus_type=corpus_info.corpus_type)
-    success = corpus.store_and_index(document, document_name, citation)
+    success = corpus.store_and_index(document, citation)
 
     current_app.logger.info(f"Memorize finished [success={success}]")
     return {"success": success}

memas/interface/corpus.py

@@ -14,6 +14,7 @@ class Citation:
     source_uri: str
     source_name: str
     description: str
+    document_name: str
 
 
 @dataclass
@@ -29,11 +30,11 @@ class Corpus(ABC):
     """
 
     def __init__(self, corpus_id: UUID, corpus_name: str):
-        self.corpus_id = corpus_id
-        self.corpus_name = corpus_name
+        self.corpus_id: UUID = corpus_id
+        self.corpus_name: str = corpus_name
 
     @abstractmethod
-    def store_and_index(self, document: str, document_name: str, citation: Citation) -> bool:
+    def store_and_index(self, document: str, citation: Citation) -> bool:
         """Store and index a "document"
 
         Args:

memas/interface/storage_driver.py

@@ -92,13 +92,12 @@ class CorpusDocumentMetadataStore(StorageDriver):
         Metadata store for storing citations and other metadata for documents within the corpus.
     """
     @abstractmethod
-    def insert_document_metadata(self, corpus_id: UUID, document_id: UUID, num_segments: int, document_name: str, citation: Citation) -> bool:
+    def insert_document_metadata(self, corpus_id: UUID, document_id: UUID, num_segments: int, citation: Citation) -> bool:
         """Inserts document metadata
 
         Args:
             corpus_id (UUID): corpus id
             document_id (UUID): document id
-            document_name (str): document name
             citation (Citation): citation object
 
         Returns:
@@ -122,6 +121,7 @@ def get_document_citation(self, corpus_id: UUID, document_id: UUID) -> Citation:
 class DocumentEntity:
     corpus_id: UUID
     document_id: UUID
+    # while strictly speaking this is metadata, this increases data readability
     document_name: str
     document: str
 

memas/storage_driver/corpus_doc_metadata.py

@@ -31,13 +31,12 @@ def init(self):
     def first_init(self):
         self.init()
 
-    def insert_document_metadata(self, corpus_id: UUID, document_id: UUID, num_segments: int, document_name: str, citation: Citation) -> bool:
+    def insert_document_metadata(self, corpus_id: UUID, document_id: UUID, num_segments: int, citation: Citation) -> bool:
         """Inserts document metadata
 
         Args:
             corpus_id (UUID): corpus id
             document_id (UUID): document id
-            document_name (str): document name
             citation (Citation): citation object
 
         Returns:
@@ -47,7 +46,7 @@ def insert_document_metadata(self, corpus_id: UUID, document_id: UUID, num_segme
 
         DocumentMetadata.create(corpus_id=corpus_id,
                                 document_id=document_id,
-                                document_name=document_name,
+                                document_name=citation.document_name,
                                 source_name=citation.source_name,
                                 source_uri=citation.source_uri,
                                 description=citation.description,
@@ -70,7 +69,8 @@ def get_document_citation(self, corpus_id: UUID, document_id: UUID) -> Citation:
             corpus_id=corpus_id, document_id=document_id)
         return Citation(source_uri=result.source_uri,
                         source_name=result.source_name,
-                        description=result.description)
+                        description=result.description,
+                        document_name=result.document_name)
 
     def get_document_segment_count(self, corpus_id: UUID, document_id: UUID) -> int:
         """Retrieves the number of segments a stored document was split into 

memas/storage_driver/corpus_vector_store.py

@@ -41,7 +41,7 @@
                 max_length=32, is_partition_key=True),
     FieldSchema(name=DOCUMENT_NAME, dtype=DataType.VARCHAR,
                 max_length=256),
-    FieldSchema(name="text_preview", dtype=DataType.VARCHAR, max_length=MAX_TEXT_LENGTH),
+    FieldSchema(name=TEXT_PREVIEW, dtype=DataType.VARCHAR, max_length=MAX_TEXT_LENGTH),
     FieldSchema(name=EMBEDDING_FIELD, dtype=DataType.FLOAT_VECTOR, dim=USE_VECTOR_DIMENSION),
     FieldSchema(name=START_FIELD, dtype=DataType.INT64),
     FieldSchema(name=END_FIELD, dtype=DataType.INT64),