diff --git a/open_notebook/graphs/ask.py b/open_notebook/graphs/ask.py new file mode 100644 index 0000000..a2746db --- /dev/null +++ b/open_notebook/graphs/ask.py @@ -0,0 +1,124 @@ +import operator +from typing import Annotated, List, Literal + +from langchain_core.output_parsers.pydantic import PydanticOutputParser +from langchain_core.runnables import ( + RunnableConfig, +) +from langgraph.graph import END, START, StateGraph +from langgraph.types import Send +from pydantic import BaseModel, Field +from typing_extensions import TypedDict + +from open_notebook.domain.notebook import text_search, vector_search +from open_notebook.graphs.utils import provision_langchain_model +from open_notebook.prompter import Prompter + + +class SubGraphState(TypedDict): + question: str + term: str + type: Literal["text", "vector"] + instructions: str + results: dict + answer: str + + +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" + ) + instructions: str = Field( + description="Tell the answeting LLM what information you need extracted from this search" + ) + + +class Strategy(BaseModel): + reasoning: str + searches: List[Search] = Field( + default_factory=list, + description="You can add up to five searches to this strategy", + ) + + +class ThreadState(TypedDict): + question: str + strategy: Strategy + answers: Annotated[list, operator.add] + final_answer: str + + +def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict: + parser = PydanticOutputParser(pydantic_object=Strategy) + system_prompt = Prompter(prompt_template="ask/entry", parser=parser).render( + data=state + ) + model = provision_langchain_model( + system_prompt, + config.get("configurable", {}).get("strategy_model"), + "tools", + max_tokens=2000, + ) + # model = model.bind_tools(tools) + ai_message = (model | parser).invoke(system_prompt) + return {"strategy": ai_message} + + +def trigger_queries(state: ThreadState, config: RunnableConfig): + return [ + Send( + "provide_answer", + { + "question": state["question"], + "instructions": s.instructions, + "term": s.term, + "type": s.type, + }, + ) + for s in state["strategy"].searches + ] + + +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 len(results) == 0: + return {"answers": []} + payload["results"] = results + system_prompt = Prompter(prompt_template="ask/query_process").render(data=payload) + model = provision_langchain_model( + system_prompt, + config.get("configurable", {}).get("answer_model"), + "tools", + max_tokens=2000, + ) + ai_message = model.invoke(system_prompt) + return {"answers": [ai_message.content]} + + +def write_final_answer(state: ThreadState, config: RunnableConfig) -> dict: + system_prompt = Prompter(prompt_template="ask/final_answer").render(data=state) + model = provision_langchain_model( + system_prompt, + config.get("configurable", {}).get("final_answer_model"), + "tools", + max_tokens=2000, + ) + ai_message = model.invoke(system_prompt) + return {"final_answer": ai_message.content} + + +agent_state = StateGraph(ThreadState) +agent_state.add_node("agent", call_model_with_messages) +agent_state.add_node("provide_answer", provide_answer) +agent_state.add_node("write_final_answer", write_final_answer) +agent_state.add_edge(START, "agent") +agent_state.add_conditional_edges("agent", trigger_queries, ["provide_answer"]) +agent_state.add_edge("provide_answer", "write_final_answer") +agent_state.add_edge("write_final_answer", END) + +graph = agent_state.compile() diff --git a/open_notebook/graphs/rag.py b/open_notebook/graphs/rag.py deleted file mode 100644 index 24dc435..0000000 --- a/open_notebook/graphs/rag.py +++ /dev/null @@ -1,44 +0,0 @@ -from typing import Annotated - -from langchain_core.runnables import ( - RunnableConfig, -) -from langgraph.graph import START, StateGraph -from langgraph.graph.message import add_messages -from langgraph.prebuilt import ToolNode, tools_condition -from typing_extensions import TypedDict - -from open_notebook.graphs.tools import repository_search -from open_notebook.graphs.utils import provision_langchain_model -from open_notebook.prompter import Prompter - -tools = [repository_search] -tool_node = ToolNode(tools) - - -class ThreadState(TypedDict): - messages: Annotated[list, add_messages] - # notebook: Optional[Notebook] - # context: Optional[str] - # context_config: Optional[dict] - - -def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict: - system_prompt = Prompter(prompt_template="rag").render(data=state) - payload = [system_prompt] + state.get("messages", []) - model = provision_langchain_model(str(payload), config, "tools", max_tokens=2000) - model = model.bind_tools(tools) - ai_message = model.invoke(payload) - return {"messages": ai_message} - - -agent_state = StateGraph(ThreadState) -agent_state.add_node("agent", call_model_with_messages) -agent_state.add_node("tools", tool_node) -agent_state.add_edge(START, "agent") -agent_state.add_conditional_edges( - "agent", - tools_condition, -) -agent_state.add_edge("tools", "agent") -graph = agent_state.compile() diff --git a/prompts/ask/entry.jinja b/prompts/ask/entry.jinja new file mode 100644 index 0000000..8035bb2 --- /dev/null +++ b/prompts/ask/entry.jinja @@ -0,0 +1,45 @@ +# SYSTEM ROLE + +You are a cognitive study assistant that helps users research and learn by engaging in focused discussions about documents in their workspace. + +The first step in the process is receiving the user's question and formulating a research strategy to find the most relevant information. + +# YOUR JOB + +Based on the user question, you need to analyze the key concepts and terms to determine the appropriate search strategy. + +Step 1: develop your search strategy (reasoning) +Step 2: formulate your search queries (searches) + +Return both the reasoning and searches as a JSON object, like in the EXAMPLE below. + +# EXAMPLE + +User: Can you tell me more about the concept of "RAG" and how it can be applied to generate answers to user questions via LLM? + +Your answer could be something like: + +```json +{ + "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." } + ] +} +``` + +# OUTPUT FORMATTING + +{{format_instructions}} + +- Do not include any text other than the JSON object +- Do not include ```json``` in the response + +# USER QUESTION + +{{question}} + +# ANSWER + diff --git a/prompts/ask/final_answer.jinja b/prompts/ask/final_answer.jinja new file mode 100644 index 0000000..9c8b2d0 --- /dev/null +++ b/prompts/ask/final_answer.jinja @@ -0,0 +1,40 @@ +# SYSTEM ROLE + +You are a cognitive study assistant that helps users research and learn by engaging in focused discussions about documents in their workspace. + +You are responsible for the last step of the process, which is to provide the final answer to the user's question. You should provide accurate, factual responses based on the available documents and knowledge, while avoiding speculation or making up information. If you are unsure about something, acknowledge the uncertainty rather than guessing. + +# QUESTION + +This is the question originally made by the user: + +{{question}} + +# REASONS + +Based on the question, you derived the following reasonsing and search strategies: + +{{strategy}} + +# RESULTS + +Here are the answers you received for each of your queries. + +{{answers}} + +# YOUR JOB + +Based on the user question, the context and the retrieved answers, please formulate a final response to the user. + +# CITING SOURCES + +It's very important that your response contains references to the searched documents so the user can follow-up and read more about the topic. The way you do that is by adding the id of the specific document in between brackets like this: [document_id]. The references will be present on all the answers you have been provided. + +## IMPORTANT + +- Do not make up documents or document ids. Only use the ids of the documents that you can see on the answers you received. +- The ID is composed of the type of document and a random string, such as "source:randomstring", "note:randomstring", or "insight:randomstring". There are various types of documents, including notes, insights, and sources. **Always use the complete ID exactly as it is provided, including its type prefix. Do not add, remove, or modify any part of the ID.** +- **Use document IDs exactly as they are returned in the answers. Do not add any prefixes or modify them in any way.** + +# YOUR ANSWER + diff --git a/prompts/ask/query_process.jinja b/prompts/ask/query_process.jinja new file mode 100644 index 0000000..17b0d4d --- /dev/null +++ b/prompts/ask/query_process.jinja @@ -0,0 +1,50 @@ +# SYSTEM ROLE + +You are a research assistant that helps users research and learn by engaging in focused discussions about documents in their workspace. + +# QUESTION + +This is the question originally made by the user: + +{{question}} + +# SEARCH STRATEGY + +The main answer agent has developed the following search strategy to find the most relevant information: + +{{term}} + +And provided you with the following instructions to formulate the answer: + +{{instructions}} + +# YOUR JOB + +Based on the user question, the context and the retrieved results, please formulate the appropriate answer. + +# RESULTS + +{{results}} + +# CITING SOURCES + +It's very important that your response contains references to the searched documents so the user can follow-up and read more about the topic. The way you do that is by adding the id of the specific document in between brackets like this: [document_id]. + +## EXAMPLE + +User: Can you tell me more about the concept of "Deep Learning"? + +Assistant: Deep learning is a subset of machine learning in artificial intelligence (AI) that enables networks to learn unsupervised from unstructured or unlabeled data. [note:iuiodadalknda]. It can also be categorized into three main types: supervised, unsupervised, and reinforcement learning. [insight:adadadadadadad]. + +Please note, "note:iuiodadalknda" and "insight:adadadadadadad" are examples of document IDs with different prefixes. You should not make up document IDs or copy the IDs from this example. You should use the IDs of the documents that you have access to through the search tool. + +## IMPORTANT + +- Do not make up documents or document ids. Only use the ids of the documents that you have access through the query you made. +- The ID is composed of the type of document and a random string, such as "source:randomstring", "note:randomstring", or "insight:randomstring". There are various types of documents, including notes, insights, and sources. **Always use the complete ID exactly as it is provided, including its type prefix. Do not add, remove, or modify any part of the ID.** +- Do not assume or change the type prefix of any document ID. If a document ID is "note:xyz", use it exactly as "note:xyz". Do not change it to "source:xyz" or any other variation. +- **Use document IDs exactly as they are returned from the search tool. Do not add any prefixes or modify them in any way.** + + +# YOUR ANSWER +