integrate hybrid search, make weaviate param configurable and add cohere integration

api/ask_astro/chains/answer_question.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 
+from typing import Any
+
 from langchain import LLMChain
+from langchain.callbacks.manager import (
+    CallbackManagerForChainRun,
+)
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT
 from langchain.chains.question_answering import load_qa_chain
@@ -12,9 +17,11 @@
     SystemMessagePromptTemplate,
 )
 from langchain.retrievers import MultiQueryRetriever
+from langchain.retrievers.document_compressors import CohereRerank
+from langchain.retrievers.weaviate_hybrid_search import WeaviateHybridSearchRetriever
 
-from ask_astro.clients.weaviate_ import docsearch
-from ask_astro.config import AzureOpenAIParams
+from ask_astro.clients.weaviate_ import client
+from ask_astro.config import AzureOpenAIParams, CohereConfig, WeaviateConfig
 from ask_astro.settings import (
     CONVERSATIONAL_RETRIEVAL_LLM_CHAIN_DEPLOYMENT_NAME,
     CONVERSATIONAL_RETRIEVAL_LLM_CHAIN_TEMPERATURE,
@@ -39,11 +46,35 @@
         deployment_name=MULTI_QUERY_RETRIEVER_DEPLOYMENT_NAME,
         temperature=MULTI_QUERY_RETRIEVER_TEMPERATURE,
     ),
-    retriever=docsearch.as_retriever(),
+    retriever=WeaviateHybridSearchRetriever(
+        client=client,
+        index_name=WeaviateConfig.index_name,
+        text_key=WeaviateConfig.text_key,
+        attributes=WeaviateConfig.attributes,
+        create_schema_if_missing=WeaviateConfig.create_schema_if_missing,
+        k=WeaviateConfig.k,
+        alpha=WeaviateConfig.alpha,
+    ),
 )
 
-# Set up a ConversationalRetrievalChain to generate answers using the retriever.
-answer_question_chain = ConversationalRetrievalChain(
+
+class AskAstroCoversationalRetrievalChainChain(ConversationalRetrievalChain):
+    def _get_docs(
+        self,
+        question: str,
+        inputs: dict[str, Any],
+        *,
+        run_manager: CallbackManagerForChainRun,
+    ):
+        docs = super()._get_docs(question=question, inputs=inputs, run_manager=run_manager)
+        if CohereConfig.cohere_api_key:
+            compressor = CohereRerank(top_n=CohereConfig.top_n, user_agent="langchain")
+            return compressor.compress_documents(docs, question)
+        return docs
+
+
+# Set up a AskAstroCoversationalRetrievalChainChain to generate answers using the retriever.
+answer_question_chain = AskAstroCoversationalRetrievalChainChain(
     retriever=retriever,
     return_source_documents=True,
     question_generator=LLMChain(

api/ask_astro/config.py

@@ -60,3 +60,14 @@ class WeaviateConfig:
     index_name = os.environ.get("WEAVIATE_INDEX_NAME")
     text_key = os.environ.get("WEAVIATE_TEXT_KEY")
     attributes = os.environ.get("WEAVIATE_ATTRIBUTES", "").split(",")
+
+    k = int(os.environ.get("WEAVIATE_K", "5"))
+    alpha = float(os.environ.get("WEAVIATE_ALPHA", "0.5"))
+    create_schema_if_missing = bool(os.environ.get("WEAVIATE_CREATE_SCHEMA_IF_MISSING", "").lower() == "true")
+
+
+class CohereConfig:
+    """Contains the config variables for the Cohere API."""
+
+    cohere_api_key = os.environ.get("COHERE_API_KEY")
+    top_n = int(os.environ.get("COHERE_TOP_N", 3))