fixes

chain.py

@@ -43,41 +43,10 @@ from typing import List, Optional, Any
 import chromadb
 from langchain.vectorstores import Chroma
 
-class CustomChain(Chain, BaseModel):
-
-    vstore: Chroma
-    chain: BaseCombineDocumentsChain
-    key_word_extractor: Chain
-
-    @property
-    def input_keys(self) -> List[str]:
-        return ["question"]
-
-    @property
-    def output_keys(self) -> List[str]:
-        return ["answer"]
-
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
-        question = inputs["question"]
-        chat_history_str = _get_chat_history(inputs["chat_history"])
-        if chat_history_str:
-            new_question = self.key_word_extractor.run(
-                question=question, chat_history=chat_history_str
-            )
-        else:
-            new_question = question
-        print(new_question)
-        docs = self.vstore.similarity_search(new_question, k=10)
-        new_inputs = inputs.copy()
-        new_inputs["question"] = new_question
-        new_inputs["chat_history"] = chat_history_str
-        answer, _ = self.chain.combine_docs(docs, **new_inputs)
-        return {"answer": answer}
-
-
 def get_new_chain1(vectorstore, model_selector, k_textbox) -> Chain:
     max_tokens_dict = {'gpt-4': 2000, 'gpt-3.5-turbo': 1000}
 
+    # These templates aren't used for the moment.
     _eg_template = """## Example:
 
     Chat History:
@@ -106,8 +75,8 @@ Answer in Markdown:"""
 
     # Construct a ChatVectorDBChain with a streaming llm for combine docs
     # and a separate, non-streaming llm for question generation
-    llm = ChatOpenAI(client = None, temperature=0.7, model_name="gpt-4")
-    streaming_llm = ChatOpenAI(client = None, streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0.7, model_name="gpt-4", max_tokens=1000)
+    llm = ChatOpenAI(client = None, temperature=0.7, model_name=model_selector)
+    streaming_llm = ChatOpenAI(client = None, streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0.7, model_name=model_selector, max_tokens=1000)
 
     question_generator = LLMChain(llm=llm, prompt=CONDENSE_QUESTION_PROMPT)
     doc_chain = load_qa_chain(streaming_llm, chain_type="stuff", prompt=QA_PROMPT)
@@ -119,14 +88,4 @@ Answer in Markdown:"""
     qa = ConversationalRetrievalChain(
         retriever=retriever, memory=memory, combine_docs_chain=doc_chain, question_generator=question_generator)
 
-
-    return qa
-
-
-def _get_chat_history(chat_history: List[Tuple[str, str]]):
-    buffer = ""
-    for human_s, ai_s in chat_history:
-        human = f"Human: " + human_s
-        ai = f"Assistant: " + ai_s
-        buffer += "\n" + "\n".join([human, ai])
-    return buffer
+    return qa