make model rag work with vector only

This commit is contained in:
LUIS NOVO 2024-11-13 12:18:26 -03:00
parent e4b8fa8cc7
commit 80353a97c9
3 changed files with 17 additions and 15 deletions

View file

@ -1,5 +1,5 @@
import operator import operator
from typing import Annotated, List, Literal from typing import Annotated, List
from langchain_core.output_parsers.pydantic import PydanticOutputParser from langchain_core.output_parsers.pydantic import PydanticOutputParser
from langchain_core.runnables import ( from langchain_core.runnables import (
@ -7,10 +7,11 @@ from langchain_core.runnables import (
) )
from langgraph.graph import END, START, StateGraph from langgraph.graph import END, START, StateGraph
from langgraph.types import Send from langgraph.types import Send
from loguru import logger
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from typing_extensions import TypedDict from typing_extensions import TypedDict
from open_notebook.domain.notebook import text_search, vector_search from open_notebook.domain.notebook import vector_search
from open_notebook.graphs.utils import provision_langchain_model from open_notebook.graphs.utils import provision_langchain_model
from open_notebook.prompter import Prompter from open_notebook.prompter import Prompter
@ -18,7 +19,7 @@ from open_notebook.prompter import Prompter
class SubGraphState(TypedDict): class SubGraphState(TypedDict):
question: str question: str
term: str term: str
type: Literal["text", "vector"] # type: Literal["text", "vector"]
instructions: str instructions: str
results: dict results: dict
answer: str answer: str
@ -26,9 +27,9 @@ class SubGraphState(TypedDict):
class Search(BaseModel): class Search(BaseModel):
term: str term: str
type: Literal["text", "vector"] = Field( # type: Literal["text", "vector"] = Field(
description="The type of search. Use 'text' for keyword search and 'vector' for semantic search. If you are using text, search always for a single word" # description="The type of search. Use 'text' for keyword search and 'vector' for semantic search. If you are using text, search always for a single word"
) # )
instructions: str = Field( instructions: str = Field(
description="Tell the answeting LLM what information you need extracted from this search" description="Tell the answeting LLM what information you need extracted from this search"
) )
@ -62,6 +63,7 @@ async def call_model_with_messages(state: ThreadState, config: RunnableConfig) -
) )
# model = model.bind_tools(tools) # model = model.bind_tools(tools)
ai_message = (model | parser).invoke(system_prompt) ai_message = (model | parser).invoke(system_prompt)
logger.debug(ai_message)
return {"strategy": ai_message} return {"strategy": ai_message}
@ -73,7 +75,7 @@ async def trigger_queries(state: ThreadState, config: RunnableConfig):
"question": state["question"], "question": state["question"],
"instructions": s.instructions, "instructions": s.instructions,
"term": s.term, "term": s.term,
"type": s.type, # "type": s.type,
}, },
) )
for s in state["strategy"].searches for s in state["strategy"].searches
@ -82,10 +84,10 @@ async def trigger_queries(state: ThreadState, config: RunnableConfig):
async def provide_answer(state: SubGraphState, config: RunnableConfig) -> dict: async def provide_answer(state: SubGraphState, config: RunnableConfig) -> dict:
payload = state payload = state
if state["type"] == "text": # if state["type"] == "text":
results = text_search(state["term"], 10, True, True) # results = text_search(state["term"], 10, True, True)
else: # else:
results = vector_search(state["term"], 10, True, True) results = vector_search(state["term"], 10, True, True)
if len(results) == 0: if len(results) == 0:
return {"answers": []} return {"answers": []}
payload["results"] = results payload["results"] = results

View file

@ -83,7 +83,7 @@ with ask_tab:
f"Agent Strategy: {chunk['agent']['strategy'].reasoning}" f"Agent Strategy: {chunk['agent']['strategy'].reasoning}"
): ):
for search in chunk["agent"]["strategy"].searches: for search in chunk["agent"]["strategy"].searches:
st.markdown(f"**{search.type} - {search.term}**") st.markdown(f"Searched for: **{search.term}**")
st.markdown(f"Instructions: {search.instructions}") st.markdown(f"Instructions: {search.instructions}")
elif "provide_answer" in chunk: elif "provide_answer" in chunk:
for answer in chunk["provide_answer"]["answers"]: for answer in chunk["provide_answer"]["answers"]:

View file

@ -23,9 +23,9 @@ Your answer could be something like:
{ {
"reasoning": "The user is asking about the concept of RAG and its application in generating answers to user questions via LLM. I should search for documents related to RAG, retrieval augmented generation, and vector search to provide a comprehensive response.", "reasoning": "The user is asking about the concept of RAG and its application in generating answers to user questions via LLM. I should search for documents related to RAG, retrieval augmented generation, and vector search to provide a comprehensive response.",
"searches": [ "searches": [
{ "type": "text", "term": "RAG", "instructions": "Describe the concept and utility of RAG." }, { "term": "RAG", "instructions": "Describe the concept and utility of RAG." },
{ "type": "vector", "term": "Retrieval Augmented Generation", "instructions": "Describe the concept and utility of RAG." }, { "term": "Retrieval Augmented Generation", "instructions": "Describe the concept and utility of RAG." },
{ "type": "vector", "term": "Vector Search", "instructions": "Describe how RAG utilizes vector search." } { "term": "Vector Search", "instructions": "Describe how RAG utilizes vector search." }
] ]
} }
``` ```