latest fixes

app.py

@@ -1,18 +1,23 @@
 # chat-pykg/app.py
 import datetime
 import os
-import gradio as gr
+import random
+import shutil
+import string
+
 import chromadb
+import gradio as gr
 from chromadb.config import Settings
-# logging.basicConfig(stream=sys.stdout, level=logging.INFO)
-# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
-from langchain.vectorstores import Chroma
 from langchain.docstore.document import Document
-import shutil
-import random, string
+from langchain.embeddings import HuggingFaceEmbeddings, OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+
 from chain import get_new_chain1
 from ingest import ingest_docs
 
+# logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
+
 def randomword(length):
     letters = string.ascii_lowercase
     return ''.join(random.choice(letters) for i in range(length))
@@ -23,22 +28,21 @@ def change_tab():
 def merge_collections(collection_load_names, vs_state): 
     merged_documents = [] 
     merged_embeddings = []
-    client = chromadb.Client(Settings(
-        chroma_db_impl="duckdb+parquet",
-        persist_directory=".persisted_data" # Optional, defaults to .chromadb/ in the current directory
-    ))
-    
     for collection_name in collection_load_names: 
-        collection_name = collection_name
+        chroma_obj_get = chromadb.Client(Settings(
+            chroma_db_impl="duckdb+parquet",
+            persist_directory=".persisted_data",
+            anonymized_telemetry = True
+        ))
         if collection_name == '': 
             continue
-        collection = client.get_collection(collection_name)
-        collection = collection.get(include=["metadatas", "documents", "embeddings"])
+        collection_obj = chroma_obj_get.get_collection(collection_name, embedding_function=HuggingFaceEmbeddings())
+        collection = collection_obj.get(include=["metadatas", "documents", "embeddings"])
         for i in range(len(collection['documents'])):
             merged_documents.append(Document(page_content=collection['documents'][i], metadata = collection['metadatas'][i]))
             merged_embeddings.append(collection['embeddings'][i])
-    merged_collection_name = "merged_collection" 
-    merged_vectorstore = Chroma.from_documents(documents=merged_documents, embeddings=merged_embeddings, collection_name=merged_collection_name) 
+    merged_vectorstore = Chroma(collection_name="temp", embedding_function=HuggingFaceEmbeddings())
+    merged_vectorstore.add_documents(documents=merged_documents, embeddings=merged_embeddings)
     return merged_vectorstore
 
 def set_chain_up(openai_api_key, model_selector, k_textbox, max_tokens_textbox, vectorstore, agent):
@@ -66,9 +70,12 @@ def delete_collection(all_collections_state, collections_viewer):
         persist_directory=".persisted_data" # Optional, defaults to .chromadb/ in the current directory
     ))
     for collection in collections_viewer:
-        client.delete_collection(collection)
-        all_collections_state.remove(collection)
-        collections_viewer.remove(collection)
+        try:
+            client.delete_collection(collection)
+            all_collections_state.remove(collection)
+            collections_viewer.remove(collection)
+        except:
+            continue
     return all_collections_state, collections_viewer
 
 def delete_all_collections(all_collections_state):
@@ -91,6 +98,9 @@ def destroy_agent(agent):
     agent = None
     return agent
 
+def clear_chat(chatbot, history):
+    return [], []
+
 def chat(inp, history, agent):
     history = history or []
     if type(agent) == str:
@@ -144,15 +154,16 @@ with block:
                     show_label=True,
                     lines=1,
                 )
-                max_tokens_textbox.value="2000"
+                max_tokens_textbox.value="1000"
             chatbot = gr.Chatbot()
             with gr.Row():
+                clear_btn = gr.Button("Clear Chat", variant="secondary").style(full_width=False)
                 message = gr.Textbox(
                     label="What's your question?",
                     placeholder="What is this code?",
                     lines=1,
                 )
-                submit = gr.Button(value="Send", variant="secondary").style(full_width=False)
+                submit = gr.Button(value="Send").style(full_width=False)
             gr.Examples(
                 examples=[
                     "What does this code do?",
@@ -207,14 +218,14 @@ with block:
         chat_state = gr.State()
 
         submit.click(set_chain_up, inputs=[openai_api_key_textbox, model_selector, k_textbox, max_tokens_textbox, vs_state, agent_state], outputs=[agent_state]).then(chat, inputs=[message, history_state, agent_state], outputs=[chatbot, history_state])
-        message.submit(chat, inputs=[message, history_state, agent_state], outputs=[chatbot, history_state])
+        message.submit(set_chain_up, inputs=[openai_api_key_textbox, model_selector, k_textbox, max_tokens_textbox, vs_state, agent_state], outputs=[agent_state]).then(chat, inputs=[message, history_state, agent_state], outputs=[chatbot, history_state])
 
         load_collections_button.click(merge_collections, inputs=[collections_viewer, vs_state], outputs=[vs_state])#.then(change_tab, None, tabs) #.then(set_chain_up, inputs=[openai_api_key_textbox, model_selector, k_textbox, max_tokens_textbox, vs_state, agent_state], outputs=[agent_state])
         make_collections_button.click(ingest_docs, inputs=[all_collections_state, all_collections_to_get, chunk_size_textbox, chunk_overlap_textbox], outputs=[all_collections_state], show_progress=True).then(update_checkboxgroup, inputs = [all_collections_state], outputs = [collections_viewer])
         delete_collections_button.click(delete_collection, inputs=[all_collections_state, collections_viewer], outputs=[all_collections_state, collections_viewer]).then(update_checkboxgroup, inputs = [all_collections_state], outputs = [collections_viewer])
         delete_all_collections_button.click(delete_all_collections, inputs=[all_collections_state], outputs=[all_collections_state]).then(update_checkboxgroup, inputs = [all_collections_state], outputs = [collections_viewer])
         get_all_collection_names_button.click(list_collections, inputs=[all_collections_state], outputs=[all_collections_state]).then(update_checkboxgroup, inputs = [all_collections_state], outputs = [collections_viewer])
-        
+        clear_btn.click(clear_chat, inputs = [chatbot, history_state], outputs = [chatbot, history_state])
         # Whenever chain parameters change, destroy the agent. 
         input_list = [openai_api_key_textbox, model_selector, k_textbox, max_tokens_textbox]
         output_list = [agent_state]

chain.py

@@ -1,8 +1,6 @@
 # chat-pykg/chain.py
 
 from langchain.chains.base import Chain
-# logging.basicConfig(stream=sys.stdout, level=logging.INFO)
-# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
 from langchain import HuggingFaceHub
 from langchain.chains.question_answering import load_qa_chain
 from langchain.chat_models import ChatOpenAI
@@ -13,8 +11,22 @@ from langchain.callbacks.base import CallbackManager
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT, QA_PROMPT
 
+# logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
+
 def get_new_chain1(vectorstore, model_selector, k_textbox, max_tokens_textbox) -> Chain:
 
+    template = """You are called chat-pykg and are an AI assistant coded in python using langchain and gradio. You are very helpful for answering questions about various open source libraries.
+                You are given the following extracted parts of code and a question. Provide a conversational answer to the question.
+                Do NOT make up any hyperlinks that are not in the code.
+                If you don't know the answer, just say that you don't know, don't try to make up an answer.
+                If the question is not about the package documentation, politely inform them that you are tuned to only answer questions about the package documentations.
+                Question: {question}
+                =========
+                {context}
+                =========
+                Answer in Markdown:"""
+    QA_PROMPT.template = template
     if model_selector in ['gpt-4', 'gpt-3.5-turbo']:
         llm = ChatOpenAI(client = None, temperature=0.7, model_name=model_selector)
         doc_chain_llm = ChatOpenAI(client = None, streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0.7, model_name=model_selector, max_tokens=int(max_tokens_textbox))
@@ -26,7 +38,7 @@ def get_new_chain1(vectorstore, model_selector, k_textbox, max_tokens_textbox) -
 
     # memory = ConversationKGMemory(llm=llm, input_key="question", output_key="answer")
     memory = ConversationBufferWindowMemory(input_key="question", output_key="answer", k=5)
-    retriever = vectorstore.as_retriever()
+    retriever = vectorstore.as_retriever(search_type="similarity")
     if len(k_textbox) != 0:
         retriever.search_kwargs = {"k": int(k_textbox)}
     else:

ingest.py

@@ -9,6 +9,9 @@ from langchain.vectorstores import Chroma
 import shutil
 from pathlib import Path
 import subprocess
+import chromadb
+from chromadb.config import Settings
+import chromadb.utils.embedding_functions as ef
 
 # class CachedChroma(Chroma, ABC):
 #     """
@@ -72,7 +75,6 @@ def get_text(content):
 def ingest_docs(all_collections_state, urls, chunk_size, chunk_overlap):
     """Get documents from web pages."""
     all_docs = []
-
     folders=[]
     documents = []                    
     shutil.rmtree('downloaded/', ignore_errors=True)
@@ -157,13 +159,12 @@ def ingest_docs(all_collections_state, urls, chunk_size, chunk_overlap):
             # else:
             #     documents += text_splitter.split_documents(docs_by_ext[ext] 
         all_docs += documents
-        embeddings = HuggingFaceEmbeddings()
         if 'downloaded/' in folder:
             folder = '-'.join(folder.split('/')[1:])
         if folder == '.':
             folder = 'chat-pykg'
-        vectorstore = Chroma.from_documents(persist_directory=".persisted_data", documents=documents, embedding=embeddings, collection_name=folder)
-        vectorstore.persist()
+        collection = Chroma.from_documents(documents=documents, collection_name=folder, embedding=HuggingFaceEmbeddings(), persist_directory=".persisted_data")
+        collection.persist()
         all_collections_state.append(folder)
     return all_collections_state
     # embeddings = HuggingFaceEmbeddings()