diff --git a/api/main.py b/api/main.py index a94fb62..abe30a0 100644 --- a/api/main.py +++ b/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)} diff --git a/api/src/stampy_chat/chat.py b/api/src/stampy_chat/chat.py index 7144d40..276bb45 100644 --- a/api/src/stampy_chat/chat.py +++ b/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 + + diff --git a/api/src/stampy_chat/citations.py b/api/src/stampy_chat/citations.py index 5221ef6..ac4a991 100644 --- a/api/src/stampy_chat/citations.py +++ b/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))