pydanticai_research_agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9fd361f8-acb7-4e11-b6e2-f01d13e257d1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Building a Research Assistant using PydanticAI\n",
    "\n",
    "The goal of this project is to build a *research assistant (Research + Summarization)* that helps users explore a research topic by iteratively:  \n",
    "\n",
    "1. Generating search queries based on the user’s input and performing web searches.  \n",
    "2. Extracting and summarizing relevant information from the search results.  \n",
    "3. Organizing and updating the collected information to maintain the state of the assistant.  \n",
    "4. Delivering a comprehensive research report to the user, complete with cited sources.  \n",
    "\n",
    "This project takes inspiration from the LangGraph tutorial [available here](https://github.com/langchain-ai/research-rabbit/tree/main). However, it simplifies the implementation by eliminating unnecessary dependencies and complexities. It also focuses on providing a more efficient and streamlined solution using the PydanticAI framework.\n",
    "\n",
    "## Here's the architecture of the system\n",
    "![research-assistant](images/research-rabbit.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "f5fb98a4-5705-4220-b59e-2e03644553e0",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pydantic_ai import Agent, RunContext, Tool\n",
    "from pydantic_ai.models.openai import OpenAIModel\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import display, Markdown\n",
    "from pydantic import BaseModel\n",
    "from dataclasses import dataclass, field\n",
    "from typing import List\n",
    "from litellm import completion\n",
    "from tavily import TavilyClient\n",
    "import json\n",
    "import litellm\n",
    "import os\n",
    "import nest_asyncio\n",
    "\n",
    "# import logging\n",
    "# logging.basicConfig(format='%(asctime)s %(message)s')\n",
    "# logger = logging.getLogger()\n",
    "# logger.setLevel(logging.INFO)\n",
    "\n",
    "\n",
    "# Because we run the code in Jupyter lab, but not needed in production\n",
    "nest_asyncio.apply()\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "tavily_client = TavilyClient()\n",
    "litellm.set_verbose=False\n",
    "\n",
    "MAX_WEB_SEARCH_LOOPS = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "49e98d70-8fcc-4269-99a8-7b1a58265207",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_writer_system_prompt = \"\"\"Your goal is to generate targeted web search query.\n",
    "\n",
    "The query will gather information related to a specific topic.\n",
    "\n",
    "Topic:\n",
    "{research_topic}\n",
    "\n",
    "Return your query as a JSON object:\n",
    "{{\n",
    "    \"query\": \"string\",\n",
    "    \"aspect\": \"string\",\n",
    "    \"rationale\": \"string\"\n",
    "}}\n",
    "\"\"\"\n",
    "\n",
    "summarizer_system_prompt = \"\"\"Your goal is to generate a high-quality summary of the web search results.\n",
    "\n",
    "When EXTENDING an existing summary:\n",
    "1. Seamlessly integrate new information without repeating what's already covered\n",
    "2. Maintain consistency with the existing content's style and depth\n",
    "3. Only add new, non-redundant information\n",
    "4. Ensure smooth transitions between existing and new content\n",
    "\n",
    "When creating a NEW summary:\n",
    "1. Highlight the most relevant information from each source\n",
    "2. Provide a concise overview of the key points related to the report topic\n",
    "3. Emphasize significant findings or insights\n",
    "4. Ensure a coherent flow of information\n",
    "\n",
    "In both cases:\n",
    "- Focus on factual, objective information\n",
    "- Maintain a consistent technical depth\n",
    "- Avoid redundancy and repetition\n",
    "- DO NOT use phrases like \"based on the new results\" or \"according to additional sources\"\n",
    "- DO NOT add a preamble like \"Here is an extended summary ...\" Just directly output the summary.\n",
    "- DO NOT add a References or Works Cited section.\n",
    "\"\"\"\n",
    "\n",
    "reflection_system_prompt =  \"\"\"You are an expert research assistant analyzing a summary about {research_topic}.\n",
    "\n",
    "Your tasks:\n",
    "1. Identify knowledge gaps or areas that need deeper exploration\n",
    "2. Generate a follow-up question that would help expand your understanding\n",
    "3. Focus on technical details, implementation specifics, or emerging trends that weren't fully covered\n",
    "\n",
    "Ensure the follow-up question is self-contained and includes necessary context for web search.\n",
    "\n",
    "Return your analysis as a JSON object:\n",
    "{{ \n",
    "    \"knowledge_gap\": \"string\",\n",
    "    \"follow_up_query\": \"string\"\n",
    "}}\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "240884ca-3d8c-403e-a0f0-8ee3c2daf6a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_sources(sources):\n",
    "    \"\"\"\n",
    "    Formats a list of source dictionaries into a structured text for LLM input.\n",
    "\n",
    "    Args:\n",
    "        sources (list): A list of dictionaries containing \"title\", \"url\", \"content\", and \"score\".\n",
    "\n",
    "    Returns:\n",
    "        str: Formatted text beginning with \"Sources:\\n\\n\" and each source's details on separate lines.\n",
    "    \"\"\"\n",
    "    formatted_text = \"Sources:\\n\\n\"\n",
    "    for i, source in enumerate(sources, start=1):\n",
    "        formatted_text += (\n",
    "            f\"Source {i}:\\n\"\n",
    "            f\"Title: {source['title']}\\n\"\n",
    "            f\"Url: {source['url']}\\n\"\n",
    "            f\"Content: {source['content']}\\n\\n\"\n",
    "        )\n",
    "    return formatted_text.strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "id": "cd0d7fa5-33b8-4a6b-a116-22f6a71a47a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ResearchDeps:\n",
    "    research_topic: str = None\n",
    "    search_query: str = None\n",
    "    current_summary: str = None\n",
    "    final_summary: str = None\n",
    "    sources: List[str] = field(default_factory=list)\n",
    "    latest_web_search_result: str = None\n",
    "    research_loop_count: int = 0\n",
    "\n",
    "\n",
    "async def generate_search_query(ctx: RunContext[ResearchDeps]) -> str:\n",
    "    \"\"\" Generate a query for web search \"\"\"\n",
    "    # logger.info(\"==== CALLING generate_search_query... ====\")\n",
    "    print(\"==== CALLING generate_search_query... ====\")\n",
    "    response = completion(\n",
    "        model=\"gpt-4o-mini\",\n",
    "        messages=[{\"content\": query_writer_system_prompt.format(research_topic=ctx.deps.research_topic),\"role\": \"system\"}, {\"content\": \"Generate a query for Web search.\",\"role\": \"user\"}],\n",
    "        max_tokens=500,\n",
    "        format=\"json\"\n",
    "    \n",
    "    )\n",
    "    search_query = json.loads(response.choices[0].message.content)\n",
    "    # print(f\"====>search_query:{search_query}\")\n",
    "    ctx.deps.search_query = search_query[\"query\"]\n",
    "    return \"perform_web_search\"\n",
    "\n",
    "async def perform_web_search(ctx: RunContext[ResearchDeps]) -> str:\n",
    "    \"\"\" Do search and collect information \"\"\"\n",
    "    # logger.info(\"==== CALLING perform_web_search... ====\")\n",
    "    print(\"==== CALLING perform_web_search... ====\")\n",
    "    search_results = tavily_client.search(ctx.deps.search_query, include_raw_content=False, max_results=1)\n",
    "    search_string = format_sources(search_results[\"results\"])\n",
    "    ctx.deps.sources.extend(search_results[\"results\"])\n",
    "    ctx.deps.latest_web_search_result = search_string\n",
    "    ctx.deps.research_loop_count += 1\n",
    "    return \"summarize_sources\"\n",
    "\n",
    "\n",
    "async def summarize_sources(ctx: RunContext[ResearchDeps]) -> str:\n",
    "    \"\"\" Summarize the gathered sources \"\"\"\n",
    "    # logger.info(\"==== CALLING summarize_sources... ====\")\n",
    "    print(\"==== CALLING summarize_sources... ====\")\n",
    "    current_summary = ctx.deps.current_summary\n",
    "    most_recent_web_research = ctx.deps.latest_web_search_result\n",
    "    if current_summary:\n",
    "        user_prompt = (f\"Extend the existing summary: {current_summary}\\n\\n\"\n",
    "                       f\"Include new search results: {most_recent_web_research} \"\n",
    "                       f\"That addresses the following topic: {ctx.deps.research_topic}\")\n",
    "\n",
    "    else:\n",
    "        user_prompt = (\n",
    "            f\"Generate a summary of these search results: {most_recent_web_research} \"\n",
    "            f\"That addresses the following topic: {ctx.deps.research_topic}\"\n",
    "        )\n",
    "\n",
    "    response = completion(\n",
    "       model=\"gpt-4o-mini\",\n",
    "        messages=[\n",
    "            {\"content\": summarizer_system_prompt.format(research_topic=ctx.deps.research_topic),\"role\": \"system\"},\n",
    "            {\"content\": user_prompt,\"role\": \"user\"}\n",
    "        ],\n",
    "        max_tokens=1000,\n",
    "    )\n",
    "    ctx.deps.current_summary = response.choices[0].message.content\n",
    "    return \"reflect_on_summary\"\n",
    "\n",
    "\n",
    "async def reflect_on_summary(ctx: RunContext[ResearchDeps]) -> str:\n",
    "    \"\"\" Reflect on the summary and generate a follow-up query \"\"\"\n",
    "    # logger.info(\"==== CALLING reflect_on_summary... ====\")\n",
    "    print(\"==== CALLING reflect_on_summary... ====\\n\\n\")\n",
    "    response = response = completion(\n",
    "        model=\"gpt-4o-mini\",\n",
    "        messages=[\n",
    "            {\"content\": reflection_system_prompt.format(research_topic=ctx.deps.research_topic),\"role\": \"system\"},\n",
    "            {\"content\": f\"Identify a knowledge gap and generate a follow-up web search query based on our existing knowledge: {ctx.deps.current_summary}\",\"role\": \"user\"}\n",
    "        ],\n",
    "        max_tokens=500,\n",
    "        response_format={ \"type\": \"json_object\" }\n",
    "    )\n",
    "    follow_up_query = json.loads(response.choices[0].message.content)\n",
    "    ctx.deps.search_query = follow_up_query[\"follow_up_query\"]\n",
    "    return \"continue_or_stop_research\"\n",
    "\n",
    "async def finalize_summary(ctx: RunContext[ResearchDeps]) -> str:\n",
    "    \"\"\" Finalize the summary \"\"\"\n",
    "    print(\"==== CALLING finalize_summary... ====\")\n",
    "    all_sources = format_sources(ctx.deps.sources)\n",
    "    ctx.deps.final_summary = f\"## Summary:\\n\\n{ctx.deps.current_summary}\\n\\n{all_sources}\"\n",
    "    return f\"STOP and return this summary: {ctx.deps.final_summary}\"\n",
    "\n",
    "\n",
    "async def continue_or_stop_research(ctx: RunContext[ResearchDeps]) -> str:\n",
    "    \"\"\" Decide to continue the research or stop based on the follow-up query \"\"\"\n",
    "    print(\"==== CALLING continue_or_stop_research... ====\")\n",
    "    if ctx.deps.research_loop_count >= MAX_WEB_SEARCH_LOOPS:\n",
    "        await finalize_summary(ctx)\n",
    "        return \"finalize_summary\"\n",
    "    else:\n",
    "        return f\"Iterations so far: {ctx.deps.research_loop_count}.\\n\\ngenerate_search_query\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "id": "5f0248b5-bf13-4702-9fac-88e1b72f5b2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== CALLING generate_search_query... ====\n",
      "==== CALLING perform_web_search... ====\n",
      "==== CALLING summarize_sources... ====\n",
      "==== CALLING reflect_on_summary... ====\n",
      "\n",
      "\n",
      "==== CALLING continue_or_stop_research... ====\n",
      "==== CALLING generate_search_query... ====\n",
      "==== CALLING perform_web_search... ====\n",
      "==== CALLING summarize_sources... ====\n",
      "==== CALLING reflect_on_summary... ====\n",
      "\n",
      "\n",
      "==== CALLING continue_or_stop_research... ====\n",
      "==== CALLING finalize_summary... ====\n",
      "==== CALLING finalize_summary... ====\n"
     ]
    }
   ],
   "source": [
    "from pydantic_ai.models.gemini import GeminiModel\n",
    "model = OpenAIModel('gpt-4o-mini')\n",
    "# model = GeminiModel('gemini-1.5-flash-8b')\n",
    "default_system_prompt = \"\"\"You are a researcher. You need to use your tools and provide a research. \n",
    "You must STOP your research if you have done {max_loop} iterations.\n",
    "\"\"\"\n",
    "research_agent = Agent(model, system_prompt=default_system_prompt.format(max_loop=MAX_WEB_SEARCH_LOOPS),\n",
    "                       deps_type=ResearchDeps, tools=[Tool(generate_search_query), Tool(perform_web_search), \n",
    "                                                      Tool(summarize_sources), Tool(reflect_on_summary), \n",
    "                                                      Tool(finalize_summary), Tool(continue_or_stop_research)])\n",
    "\n",
    "topic = 'how to use SmolLM models?'\n",
    "\n",
    "research_deps = ResearchDeps(research_topic=topic)\n",
    "result = research_agent.run_sync(topic, deps=research_deps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 231,
   "id": "47935c8e-034d-44eb-b589-a2eddc080c8d",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "## Summary:\n",
       "\n",
       "The article provides a comprehensive guide on setting up and utilizing the SmolLM2 lightweight language model, which offers functionalities for various tasks, including text rewriting, summarization, and function calling. It details the model's availability in three sizes and outlines the process for evaluation and effective use in applications. This resource is beneficial for developers looking to incorporate lightweight language models into their projects, ensuring ease of integration and access to advanced language processing capabilities. \n",
       "\n",
       "Additionally, a step-by-step guide further clarifies the setup, evaluation, and application of SmolLM2. It emphasizes the user-friendly approach to employing this model, addressing the practical stages of integration into different workflows. The guide also underscores how to maximize the effectiveness of SmolLM2 for specific tasks, making it a valuable tool for developers interested in creating their own lightweight language model applications. This process not only enhances understanding of model usage but also expands its accessibility for a wider range of projects.\n",
       "\n",
       "### Sources:\n",
       "\n",
       "1. **Title**: Building Your Own Lightweight Language Model Applications with ... - Medium  \n",
       "   **Url**: [Read more](https://medium.com/@anoopjohny2000/building-your-own-lightweight-language-model-applications-with-smollm2-1-7b-instruct-b5cb43443891)  \n",
       "   **Content**: This step-by-step guide will help you set up, evaluate, and use the model for tasks such as text rewriting, summarization, and function calling. SmolLM2 is available in three sizes.\n",
       "\n",
       "2. **Title**: Building Your Own Lightweight Language Model Applications with ... - Medium  \n",
       "   **Url**: [Read more](https://medium.com/@anoopjohny2000/building-your-own-lightweight-language-model-applications-with-smollm2-1-7b-instruct-b5cb43443891)  \n",
       "   **Content**: This step-by-step guide will help you set up, evaluate, and use the model for tasks such as text rewriting, summarization, and function calling. SmolLM2 is available in three sizes."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 231,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import display, Markdown\n",
    "\n",
    "Markdown(result.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 234,
   "id": "4f5ab145-5a68-4d50-9fe7-4442bb300fe9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "## Summary:\n",
       "\n",
       "The article provides a comprehensive guide on setting up and utilizing the SmolLM2 lightweight language model, which offers functionalities for various tasks, including text rewriting, summarization, and function calling. It details the model's availability in three sizes and outlines the process for evaluation and effective use in applications. This resource is beneficial for developers looking to incorporate lightweight language models into their projects, ensuring ease of integration and access to advanced language processing capabilities. \n",
       "\n",
       "Additionally, a step-by-step guide further clarifies the setup, evaluation, and application of SmolLM2. It emphasizes the user-friendly approach to employing this model, addressing the practical stages of integration into different workflows. The guide also underscores how to maximize the effectiveness of SmolLM2 for specific tasks, making it a valuable tool for developers interested in creating their own lightweight language model applications. This process not only enhances understanding of model usage but also expands its accessibility for a wider range of projects.\n",
       "\n",
       "### Sources:\n",
       "\n",
       "1. **Title**: Building Your Own Lightweight Language Model Applications with ... - Medium  \n",
       "   **Url**: [Read more](https://medium.com/@anoopjohny2000/building-your-own-lightweight-language-model-applications-with-smollm2-1-7b-instruct-b5cb43443891)  \n",
       "   **Content**: This step-by-step guide will help you set up, evaluate, and use the model for tasks such as text rewriting, summarization, and function calling. SmolLM2 is available in three sizes.\n",
       "\n",
       "2. **Title**: Building Your Own Lightweight Language Model Applications with ... - Medium  \n",
       "   **Url**: [Read more](https://medium.com/@anoopjohny2000/building-your-own-lightweight-language-model-applications-with-smollm2-1-7b-instruct-b5cb43443891)  \n",
       "   **Content**: This step-by-step guide will help you set up, evaluate, and use the model for tasks such as text rewriting, summarization, and function calling. SmolLM2 is available in three sizes."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 234,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Markdown(result.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 237,
   "id": "72900ee0-c4d0-4dc0-a816-3c04dee3a4c9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "## Summary:\n",
       "\n",
       "The article provides a comprehensive guide on setting up and utilizing the SmolLM2 lightweight language model, which offers functionalities for various tasks, including text rewriting, summarization, and function calling. It details the model's availability in three sizes and outlines the process for evaluation and effective use in applications. This resource is beneficial for developers looking to incorporate lightweight language models into their projects, ensuring ease of integration and access to advanced language processing capabilities. \n",
       "\n",
       "Additionally, a step-by-step guide further clarifies the setup, evaluation, and application of SmolLM2. It emphasizes the user-friendly approach to employing this model, addressing the practical stages of integration into different workflows. The guide also underscores how to maximize the effectiveness of SmolLM2 for specific tasks, making it a valuable tool for developers interested in creating their own lightweight language model applications. This process not only enhances understanding of model usage but also expands its accessibility for a wider range of projects.\n",
       "\n",
       "Sources:\n",
       "\n",
       "Source 1:\n",
       "Title: Building Your Own Lightweight Language Model Applications with ... - Medium\n",
       "Url: https://medium.com/@anoopjohny2000/building-your-own-lightweight-language-model-applications-with-smollm2-1-7b-instruct-b5cb43443891\n",
       "Content: This step-by-step guide will help you set up, evaluate, and use the model for tasks such as text rewriting, summarization, and function calling. Model Overview SmolLM2 is available in three sizes\n",
       "\n",
       "Source 2:\n",
       "Title: Building Your Own Lightweight Language Model Applications with ... - Medium\n",
       "Url: https://medium.com/@anoopjohny2000/building-your-own-lightweight-language-model-applications-with-smollm2-1-7b-instruct-b5cb43443891\n",
       "Content: This step-by-step guide will help you set up, evaluate, and use the model for tasks such as text rewriting, summarization, and function calling. Model Overview SmolLM2 is available in three sizes"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 237,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Markdown(research_deps.final_summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5faea41f-5960-4dd9-891a-e5c841be354c",
   "metadata": {},
   "source": [
    "## Notes:\n",
    "- This is the basic version and there is a lot of room for impimprovement\n",
    "- When used with small LLM models, it's less effective in generating detailed and accurate search queries and summaries, but it can still provide a useful starting point for further research."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e41452f-3212-4629-a101-2f20b7c4ee6a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}