transform query

api/main.py

@@ -12,7 +12,7 @@ from flask_cors import CORS, cross_origin
 from stampy_chat import logging
 from stampy_chat.env import PINECONE_INDEX, FLASK_PORT, LANGCHAIN_API_KEY, LANGCHAIN_PROJECT
 from stampy_chat.settings import Settings
-from stampy_chat.chat import run_query
+from stampy_chat.chat import run_query, transform_query
 from stampy_chat.callbacks import stream_callback
 from stampy_chat.citations import get_top_k_blocks
 from stampy_chat.db.session import make_session
@@ -40,6 +40,12 @@ def semantic():
 
     query = request.json['query']
     k = request.json.get('k', 20)
+
+    # try elastic and other searches
+    n = transform_query(query)
+    logging.info(f'query=`{query}`, new_query=`{n["new_query"]}` example=`{n["example"]}`')
+    query = f'{query}. {n["new_query"]}. {n["example"]}'
+
     return jsonify(get_top_k_blocks(query, k))
 
 
@@ -61,6 +67,12 @@ def chat():
         query = history[-1].get('content')
         history = history[:-1]
 
+    # TODO: try elastic and other searches
+    # TODO ideally separate searches for example and both queries
+    n = transform_query(query)
+    logging.info(f'query=`{query}`, new_query=`{n["new_query"]}` example=`{n["example"]}`')
+    query = f'{query}. {n["new_query"]}. {n["example"]}'
+
     def formatter(item):
         if isinstance(item, Exception):
             item = {'state': 'error', 'error': str(item)}

api/src/stampy_chat/chat.py

@@ -163,6 +163,24 @@ class ModeratedChatPrompt(ChatPromptTemplate):
 def get_model(**kwargs):
     return ChatOpenAI(openai_api_key=OPENAI_API_KEY, **kwargs)
 
+def transform_query(query: str) -> str:
+    """
+    Transform the query into a new query and an example answer for better similarity search.
+    """
+    llm=get_model(model=COMPLETIONS_MODEL)
+    prompt_template = ChatPromptTemplate.from_messages([
+        ChatMessagePromptTemplate.from_template(template="Please draft a search query that will find documents to answer the following question: {query}", role='user'),
+    ])
+    chain = LLMChain(llm=llm, prompt=prompt_template)
+    new_query = chain.invoke({"query": query})['text']
+
+    prompt_template = ChatPromptTemplate.from_messages([
+        ChatMessagePromptTemplate.from_template(template="Please draft an concrete and concise example answer to the following question: {query}", role='user'),
+    ])
+    chain = LLMChain(llm=llm, prompt=prompt_template)
+    example_answer = chain.invoke({"query": query})['text']
+    return dict(new_query=new_query, example_answer=example_answer)
+
 
 class LLMInputsChain(LLMChain):
 
@@ -310,3 +328,5 @@ def run_query(session_id: str, query: str, history: List[Dict], settings: Settin
         callback({'state': 'done'})
         callback(None)  # make sure the callback handler know that things have ended
     return result
+
+

api/src/stampy_chat/citations.py

@@ -11,11 +11,10 @@ from langchain.prompts import (
 from langchain.pydantic_v1 import Extra
 from langchain.vectorstores import Pinecone
 
-from stampy_chat.env import PINECONE_INDEX, PINECONE_NAMESPACE, OPENAI_API_KEY, REMOTE_CHAT_INSTANCE
+from stampy_chat.env import PINECONE_INDEX, PINECONE_NAMESPACE, OPENAI_API_KEY, REMOTE_CHAT_INSTANCE, COMPLETIONS_MODEL
 from stampy_chat.callbacks import StampyCallbackHandler
 
 
-
 class RemoteVectorStore(VectorStore):
     """Make a wrapper around the deployed semantic search.
 
@@ -139,6 +138,7 @@ def format_block(block) -> Dict[str, Any]:
     }
 
 
+
 def get_top_k_blocks(query, k):
     blocks = make_example_selector(k=k).select_examples({'query': query})
     return list(map(format_block, blocks))