adding guidance rag agent

chat.ipynb

@@ -20,7 +20,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "73bd3df7",
    "metadata": {},
    "outputs": [
@@ -40,7 +40,7 @@
     "# Load index for retrieval\n",
     "storage_context = StorageContext.from_defaults(persist_dir=\"arv_metadata\")\n",
     "index = load_index_from_storage(storage_context)\n",
-    "retriever = index.as_retriever(similarity_top_k=10,\n",
+    "retriever = index.as_retriever(similarity_top_k=5,\n",
     "                                # Similarity threshold for filtering\n",
     "                                similarity_threshold=0.5)\n",
     "\n",

chatlib/guidlines_rag_agent_li.py

@@ -0,0 +1,32 @@
+from llama_index.core import StorageContext, load_index_from_storage
+from .state_types import State
+
+# Load index for retrieval
+storage_context = StorageContext.from_defaults(persist_dir="arv_metadata")
+index = load_index_from_storage(storage_context)
+retriever = index.as_retriever(similarity_top_k=5,
+                                # Similarity threshold for filtering
+                                similarity_threshold=0.5)
+
+def rag_retrieve(state:State) -> State:
+    """Perform RAG search of repository containing authoritative information on HIV/AIDS in Kenya.
+
+    """
+    user_prompt = state["question"]  # or whatever key holds the prompt
+    sources = retriever.retrieve(user_prompt)
+    retrieved_text = "\n\n".join([f"Source {i+1}: {source.text}" for i, source in enumerate(sources)])
+
+    return {**state, "rag_result": "RAG search results for: " + retrieved_text}
+
+if __name__ == "__main__":
+    # Test the function
+    test_state = State(
+        messages=[],
+        question="What are the first-line treatments for HIV in Kenya?",
+        rag_result="",
+        query="",
+        result="",
+        answer=""
+    )
+    updated_state = rag_retrieve(test_state)
+    print(updated_state["rag_result"])

chatlib/state_types.py

@@ -0,0 +1,11 @@
+from typing_extensions import TypedDict, Annotated
+from langchain_core.messages import AnyMessage
+from langgraph.graph.message import add_messages
+
+class State(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
+    question: str
+    rag_result: str
+    query: str
+    result: str
+    answer: str

main.py

@@ -0,0 +1,56 @@
+from dotenv import load_dotenv
+import os
+from langchain_openai import ChatOpenAI
+from langgraph.graph import MessagesState, START, StateGraph
+from langchain_core.messages import HumanMessage, SystemMessage
+from langgraph.prebuilt import tools_condition, ToolNode
+from langgraph.checkpoint.memory import MemorySaver
+memory = MemorySaver()
+
+load_dotenv("config.env")
+os.environ.get("OPENAI_API_KEY")
+os.environ.get("LANGSMITH_API_KEY")
+
+from chatlib.state_types import State
+from chatlib.guidlines_rag_agent_li import rag_retrieve
+
+tools = [rag_retrieve]
+llm = ChatOpenAI(temperature = 0.0, model="gpt-4o")
+llm_with_tools = llm.bind_tools([rag_retrieve])
+
+# System message
+sys_msg = SystemMessage(content="""
+                        You are a helpful assistant tasked with helping clinicians
+                        access information from HIV clinical guidelines.
+                        """
+                        )
+
+# Assistant Node
+def assistant(state: MessagesState):
+   return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
+
+# Graph
+builder = StateGraph(MessagesState)
+
+# Define nodes: these do the work
+builder.add_node("assistant", assistant)
+builder.add_node("tools", ToolNode(tools))
+
+# Define edges: these determine how the control flow moves
+builder.add_edge(START, "assistant")
+builder.add_conditional_edges(
+    "assistant",
+    # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
+    # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
+    tools_condition,
+)
+builder.add_edge("tools", "assistant")
+react_graph = builder.compile(checkpointer=memory)
+
+# Specify a thread
+config = {"configurable": {"thread_id": "11"}}
+
+messages = [HumanMessage(content="What are the first-line treatments for HIV in Kenya?")]
+messages = react_graph.invoke({"messages": messages}, config)
+for m in messages['messages']:
+    m.pretty_print()

requirements.txt

@@ -14,4 +14,5 @@ langgraph-cli[inmem]
 llama_index==0.12.34
 pylint
 black
-pytest
+pytest
+dotenv

sql_agent.ipynb

@@ -41,37 +41,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
-   "id": "40ddb630",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[QuerySQLDatabaseTool(description=\"Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.\", db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x707a407ab3e0>),\n",
-       " InfoSQLDatabaseTool(description='Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x707a407ab3e0>),\n",
-       " ListSQLDatabaseTool(db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x707a407ab3e0>),\n",
-       " QuerySQLCheckerTool(description='Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x707a407ab3e0>, llm=ChatOpenAI(client=<openai.resources.chat.completions.completions.Completions object at 0x707a4004e300>, async_client=<openai.resources.chat.completions.completions.AsyncCompletions object at 0x707a33d0d370>, root_client=<openai.OpenAI object at 0x707a40194050>, root_async_client=<openai.AsyncOpenAI object at 0x707a4004df70>, model_name='gpt-4o', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********')), llm_chain=LLMChain(verbose=False, prompt=PromptTemplate(input_variables=['dialect', 'query'], input_types={}, partial_variables={}, template='\\n{query}\\nDouble check the {dialect} query above for common mistakes, including:\\n- Using NOT IN with NULL values\\n- Using UNION when UNION ALL should have been used\\n- Using BETWEEN for exclusive ranges\\n- Data type mismatch in predicates\\n- Properly quoting identifiers\\n- Using the correct number of arguments for functions\\n- Casting to the correct data type\\n- Using the proper columns for joins\\n\\nIf there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.\\n\\nOutput the final SQL query only.\\n\\nSQL Query: '), llm=ChatOpenAI(client=<openai.resources.chat.completions.completions.Completions object at 0x707a4004e300>, async_client=<openai.resources.chat.completions.completions.AsyncCompletions object at 0x707a33d0d370>, root_client=<openai.OpenAI object at 0x707a40194050>, root_async_client=<openai.AsyncOpenAI object at 0x707a4004df70>, model_name='gpt-4o', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********')), output_parser=StrOutputParser(), llm_kwargs={}))]"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from langchain_community.agent_toolkits import SQLDatabaseToolkit\n",
-    "\n",
-    "toolkit = SQLDatabaseToolkit(db=db, llm=llm)\n",
-    "\n",
-    "tools = toolkit.get_tools()\n",
-    "\n",
-    "tools"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "f9c96976",
    "metadata": {},
    "outputs": [
@@ -131,8 +101,8 @@
     "    [(\"system\", system_message), (\"user\", user_prompt)]\n",
     ")\n",
     "\n",
-    "for message in query_prompt_template.messages:\n",
-    "    message.pretty_print()"
+    "# for message in query_prompt_template.messages:\n",
+    "#     message.pretty_print()"
    ]
   },
   {
@@ -349,7 +319,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 17,
    "id": "5fb11ed6",
    "metadata": {},
    "outputs": [
@@ -357,9 +327,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'assistant': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_fFrF2w9DXG12kOFn4Ox5wnGJ', 'function': {'arguments': '{\"state\":{\"messages\":[{\"content\":\"how many unique regimens are there?\",\"type\":\"human\"}],\"question\":\"how many unique regimens are there?\",\"query\":\"SELECT COUNT(DISTINCT regimen) FROM patient_records;\",\"result\":\"42\",\"answer\":\"There are 42 unique regimens.\"}}', 'name': 'sql_chain'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 71, 'prompt_tokens': 3149, 'total_tokens': 3220, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_07871e2ad8', 'id': 'chatcmpl-BkBc7T5CPXYLG1qFVcWony1uByR15', 'service_tier': 'default', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--65558f43-fd67-4e1c-83dc-068693dcda8e-0', tool_calls=[{'name': 'sql_chain', 'args': {'state': {'messages': [{'content': 'how many unique regimens are there?', 'type': 'human'}], 'question': 'how many unique regimens are there?', 'query': 'SELECT COUNT(DISTINCT regimen) FROM patient_records;', 'result': '42', 'answer': 'There are 42 unique regimens.'}}, 'id': 'call_fFrF2w9DXG12kOFn4Ox5wnGJ', 'type': 'tool_call'}], usage_metadata={'input_tokens': 3149, 'output_tokens': 71, 'total_tokens': 3220, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
-      "{'tools': {'messages': [ToolMessage(content=\"{'messages': [HumanMessage(content='how many unique regimens are there?', additional_kwargs={}, response_metadata={})], 'question': 'how many unique regimens are there?', 'query': 'SELECT COUNT(DISTINCT CurrentRegimen) AS unique_regimen_count FROM clinical_visits;', 'result': '[(12,)]', 'answer': 'There are 12 unique regimens.'}\", name='sql_chain', id='80f0822d-ee23-4c9d-a28c-359a1200a173', tool_call_id='call_fFrF2w9DXG12kOFn4Ox5wnGJ')]}}\n",
-      "{'assistant': {'messages': [AIMessage(content='There are 12 unique regimens.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 3309, 'total_tokens': 3318, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 3072}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_07871e2ad8', 'id': 'chatcmpl-BkBcBVVLMPEjPvVpK7Ap8L24bNDYH', 'service_tier': 'default', 'finish_reason': 'stop', 'logprobs': None}, id='run--4dea7411-1b2b-4f3b-b8e0-09b2407221fd-0', usage_metadata={'input_tokens': 3309, 'output_tokens': 9, 'total_tokens': 3318, 'input_token_details': {'audio': 0, 'cache_read': 3072}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n"
+      "{'assistant': {'messages': [AIMessage(content=\"You're welcome! If you have any more questions, feel free to ask.\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 16, 'prompt_tokens': 3327, 'total_tokens': 3343, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_07871e2ad8', 'id': 'chatcmpl-BkBjCRm7HUlUQZV3ntGa2Sbpz8UPJ', 'service_tier': 'default', 'finish_reason': 'stop', 'logprobs': None}, id='run--45e27132-7994-4010-91e7-cd18113d643c-0', usage_metadata={'input_tokens': 3327, 'output_tokens': 16, 'total_tokens': 3343, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n"
      ]
     }
    ],
@@ -367,7 +335,7 @@
     "# Specify a thread\n",
     "config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
     "\n",
-    "user_prompt = \"sorry, i meant to ask how many unique regimens are there?\"\n",
+    "user_prompt = \"thanks!?\"\n",
     "input_state = {\n",
     "    \"messages\": [HumanMessage(content=user_prompt)],\n",
     "    \"question\": user_prompt,\n",