From 80353a97c98e61ef96b026ebdd49f4c703640c79 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 13 Nov 2024 12:18:26 -0300
Subject: [PATCH] make model rag work with vector only

---
 open_notebook/graphs/ask.py  | 24 +++++++++++++-----------
 pages/3_🔍_Ask_and_Search.py |  2 +-
 prompts/ask/entry.jinja      |  6 +++---
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/open_notebook/graphs/ask.py b/open_notebook/graphs/ask.py
index 8201783..995864e 100644
--- a/open_notebook/graphs/ask.py
+++ b/open_notebook/graphs/ask.py
@@ -1,5 +1,5 @@
 import operator
-from typing import Annotated, List, Literal
+from typing import Annotated, List
 
 from langchain_core.output_parsers.pydantic import PydanticOutputParser
 from langchain_core.runnables import (
@@ -7,10 +7,11 @@ from langchain_core.runnables import (
 )
 from langgraph.graph import END, START, StateGraph
 from langgraph.types import Send
+from loguru import logger
 from pydantic import BaseModel, Field
 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.prompter import Prompter
 
@@ -18,7 +19,7 @@ from open_notebook.prompter import Prompter
 class SubGraphState(TypedDict):
     question: str
     term: str
-    type: Literal["text", "vector"]
+    # type: Literal["text", "vector"]
     instructions: str
     results: dict
     answer: str
@@ -26,9 +27,9 @@ class SubGraphState(TypedDict):
 
 class Search(BaseModel):
     term: str
-    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"
-    )
+    # 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"
+    # )
     instructions: str = Field(
         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)
     ai_message = (model | parser).invoke(system_prompt)
+    logger.debug(ai_message)
     return {"strategy": ai_message}
 
 
@@ -73,7 +75,7 @@ async def trigger_queries(state: ThreadState, config: RunnableConfig):
                 "question": state["question"],
                 "instructions": s.instructions,
                 "term": s.term,
-                "type": s.type,
+                # "type": s.type,
             },
         )
         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:
     payload = state
-    if state["type"] == "text":
-        results = text_search(state["term"], 10, True, True)
-    else:
-        results = vector_search(state["term"], 10, True, True)
+    # if state["type"] == "text":
+    #     results = text_search(state["term"], 10, True, True)
+    # else:
+    results = vector_search(state["term"], 10, True, True)
     if len(results) == 0:
         return {"answers": []}
     payload["results"] = results
diff --git a/pages/3_🔍_Ask_and_Search.py b/pages/3_🔍_Ask_and_Search.py
index b07e8e5..d84e404 100644
--- a/pages/3_🔍_Ask_and_Search.py
+++ b/pages/3_🔍_Ask_and_Search.py
@@ -83,7 +83,7 @@ with ask_tab:
                     f"Agent Strategy: {chunk['agent']['strategy'].reasoning}"
                 ):
                     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}")
             elif "provide_answer" in chunk:
                 for answer in chunk["provide_answer"]["answers"]:
diff --git a/prompts/ask/entry.jinja b/prompts/ask/entry.jinja
index 8035bb2..a97f03a 100644
--- a/prompts/ask/entry.jinja
+++ b/prompts/ask/entry.jinja
@@ -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.", 
     "searches": [
-        { "type": "text", "term": "RAG", "instructions": "Describe the concept and utility of RAG." },
-        { "type": "vector", "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": "RAG", "instructions": "Describe the concept and utility of RAG." },
+        { "term": "Retrieval Augmented Generation", "instructions": "Describe the concept and utility of RAG." },
+        { "term": "Vector Search", "instructions": "Describe how RAG utilizes vector search." }
     ]
 }
 ```