Some changes in chain.

doc_generator/query/__init__.py

@@ -38,7 +38,7 @@ def query(repo_config: AutodocRepoConfig, user_confg: AutodocUserConfig):
 
         print('Thinking...')
         try:
-            response = chain({'question': question, 'chat_history': chat_history})
+            response = chain.invoke({'question': question, 'chat_history': chat_history})
             chat_history.append((question, response['text']))
             print('\n\nMarkdown:\n')
             print(markdown(response['text']))

doc_generator/utils/HNSWLib.py

@@ -97,7 +97,7 @@ def from_documents(documents: List[Document], embeds: Embeddings, docstore: Opti
         hnsw.add_documents(documents)
         return hnsw
 
-    def similarity_search_by_vector(self, query: List[float], k: int) -> List:
+    def similarity_search_by_vector(self, query: List[float], k: int = 4) -> List:
         if len(query) != self.args.num_dimensions:
             raise ValueError(f"Query vector must have the same length as the number of dimensions ({self.args.num_dimensions})")
         total = self._index.element_count
@@ -107,7 +107,7 @@ def similarity_search_by_vector(self, query: List[float], k: int) -> List:
         labels, distances = self._index.knn_query(query, k)
         return [(self.docstore.search(str(label)), distance) for label, distance in zip(labels, distances)]
 
-    def similarity_search(self, query: str, k: int) -> List[Document]:
+    def similarity_search(self, query: str, k: int = 4) -> List[Document]:
         return self.similarity_search_by_vector(self._embeddings.embed_query(query), k)
 
     def save(self, directory: str):
@@ -130,15 +130,15 @@ def load(directory: str, embeddings: Embeddings):
         args = HNSWLibArgs(space=args_data['space'], num_dimensions=args_data['num_dimensions'])
         index = hnswlib.Index(space=args.space, dim=args.num_dimensions)
         index.load_index(os.path.join(directory, 'hnswlib.index'))
+        args.docstore = InMemoryDocstore()
         with open(os.path.join(directory, 'docstore.json'), 'r') as f:
             doc_data = json.load(f)
-        doc_map = {}
         for id, value in doc_data:
-            doc_map[id] = Document(
+            args.docstore.add({id: Document(
                 page_content=value['page_content'], 
                 metadata=value['metadata'],
                 type=value['type']
-            )
-        args.docstore = InMemoryDocstore([doc_map])
+            )})
+
         args.index = index
         return HNSWLib(embeddings, args)

doc_generator/utils/createChatChain.py

@@ -1,4 +1,4 @@
-from langchain.chains.conversational_retrieval.base import ChatVectorDBChain
+from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
 from langchain.chains.llm import LLMChain
 from langchain.chains.question_answering import load_qa_chain
 from langchain.prompts import PromptTemplate
@@ -87,8 +87,8 @@ def make_chain(project_name, repository_url, content_type, chat_prompt, target_a
         prompt=qa_prompt
     )
 
-    return ChatVectorDBChain(
-        vectorstore=vectorstore,
+    return ConversationalRetrievalChain(
+        retriever=vectorstore.as_retriever(),
         combine_docs_chain=doc_chain,
         question_generator=question_generator
     )