Upload 5 files

app.py

@@ -0,0 +1,348 @@
+import streamlit as st
+import pandas as pd
+from main_agent import Main_agent
+import asyncio
+import tempfile
+import os
+from PIL import Image
+from io import BytesIO
+import requests
+from spire.doc import Document, FileFormat
+import nest_asyncio
+import platform
+from dataclasses import dataclass
+
+# Apply nest_asyncio to allow nested event loops
+nest_asyncio.apply()
+
+# Create event loop
+if platform.system() == 'Windows':
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+test_data={'title': 'Sunday App and Stripe: A Payment Strategy for Restaurants',
+ 'image_url': 'https://www.menutiger.com/_next/image?url=http%3A%2F%2Fcms.menutiger.com%2Fwp-content%2Fuploads%2F2023%2F07%2Fseamless-stripe-payment-integration-for-restaurants-2.jpg&w=2048&q=75',
+ 'paragraphs': [{'title': 'Introduction',
+   'content': 'Sunday is a payment solution designed to streamline the restaurant experience. Co-founded in April 2021 by Christine de Wendel, Victor Lugger, and Tigrane Seydoux, it aims to address inefficiencies in the restaurant industry by simplifying payments through QR code technology, allowing customers to quickly access their bill, split it, add a tip, and pay in seconds. Integrated with existing POS systems and partnering with payment providers like Stripe and Checkout.com, Sunday also offers features such as digital menus with customizable options, order and pay capabilities, and tools to enhance online reputation and customer engagement for restaurants.'},
+  {'title': "Sunday's Payment Strategy",
+   'content': "Sunday's payment strategy focuses on streamlining the payment process in restaurants through its QR code-based system, enabling customers to pay quickly and efficiently. This approach not only simplifies payments but also reduces processing fees for restaurants through partnerships with companies like Stripe and Checkout.com. The company's payment solution is designed to integrate seamlessly with existing POS systems, offering additional features such as digital menus, order and pay capabilities, and tools to improve online reputation and customer engagement."},
+  {'title': 'Partnership with Stripe',
+   'content': 'On Wednesday, Sunday announced its new Order & Pay capabilities and its partnership with Stripe, which has allowed it to reduce processing fees for restaurants by an average of 0.5%. Through Stripe, Sunday processes payments via a QR code at the table, which allows customers to pay, split the bill, and add a tip quickly. Sunday is also negotiating with payment providers for better rates.'},
+  {'title': 'Stripe Integration and Fee Reduction',
+   'content': 'Sunday leverages Stripe to process payments, utilizing a QR code system for customer convenience. This partnership has enabled Sunday to reduce processing fees for restaurants by an average of 0.5%.'},
+  {'title': 'Benefits for Restaurants',
+   'content': 'Sunday offers numerous benefits for restaurants by leveraging its technology and partnerships, including Stripe. Restaurants using Sunday can experience increased efficiency through streamlined payment processes and order and pay capabilities. Sunday also helps restaurants collect valuable data and analytics, including customer insights and business analytics, which can be used to improve operations and enhance customer experience. Additionally, the integration with Stripe reduces processing fees, benefiting restaurant finances.'},
+  {'title': 'Conclusion',
+   'content': 'In conclusion, Sunday has significantly impacted restaurant payments by streamlining the payment process and reducing processing fees through its partnership with Stripe. This collaboration has enabled restaurants to enhance efficiency, improve customer experiences, and gain valuable data insights, positioning Sunday as a key player in revolutionizing the restaurant industry.'}],
+ 'table': {'data': [['App Functionality',
+    'Access bill and pay in two clicks, browse the full menu.'],
+   ['Payment Processing', 'Payments processed by Stripe.'],
+   ['Stripe Agreement Benefit',
+    'Reduced processing fees for restaurants by an average of 0.5%.'],
+   ['Restaurant Benefits',
+    'More tips, more reviews, and more data for restaurants. Increased order and average basket size.'],
+   ['Point of Sale Integration',
+    'Works seamlessly with existing POS systems.'],
+   ['Payment Method', 'Guests scan a QR code to pay.'],
+   ['Additional Features',
+    'Digital menu with customizable features (scheduling, tags, extras). Options to split the bill and tip.']],
+  'columns': ['Aspect', 'Details']},
+ 'references': 'https://sundayapp.com/2024-recap-how-sunday-revolutionized-restaurant-payments/, https://sundayapp.com/blog/, https://www.checkout.com/case-studies/sunday-brings-the-ease-of-online-payments-to-the-offline-dining-experience'}
+
+@dataclass
+class Research_paper_test:
+    test_data:dict
+
+test_paper=Research_paper_test
+test_paper.test_data=test_data
+
+
+# Initialize the agent
+agent = Main_agent()
+
+# Initialize session state
+if 'chat_history' not in st.session_state:
+    st.session_state.chat_history = []
+
+if 'table_data' not in st.session_state:
+    st.session_state.table_data = {}
+if 'quick_search_results' not in st.session_state:
+    st.session_state.quick_search_results = []
+if 'research_paper' not in st.session_state:
+    st.session_state.research_paper = {}
+
+if 'agent_memory' not in st.session_state:
+    st.session_state.agent_memory = []
+
+if 'loop' not in st.session_state:
+    st.session_state.loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(st.session_state.loop)
+
+agent.deps.deep_search_results = st.session_state.research_paper
+agent.deps.quick_search_results = st.session_state.quick_search_results
+agent.deps.table_data = st.session_state.table_data
+agent.memory.messages = st.session_state.agent_memory
+
+def paper_to_markdown(paper: dict) -> str:
+    """Convert a paper dictionary into a markdown string."""
+    markdown = []
+    
+    # Add title
+    markdown.append(f"# {paper.get('title')}")
+    markdown.append("\n\n")
+    if paper.get('image_url'):
+        markdown.append(f"![Image]({paper.get('image_url')})\n")
+        markdown.append("\n\n")
+    
+    # Add body sections
+    for section in paper.get('paragraphs', []):
+        markdown.append(f"## {section.get('title')}\n")
+        if section.get('content'):
+            markdown.append(section.get('content'))
+            markdown.append("\n")
+        markdown.append("\n\n")
+    
+    if paper.get('table'):
+        try:
+            markdown.append(pd.DataFrame(data=paper.get('table')['data'], 
+                                      columns=paper.get('table')['columns']).to_markdown())
+        except:
+            markdown.append(pd.DataFrame(paper.get('table')).to_markdown())
+        markdown.append("\n\n")
+    
+    # Add references
+    markdown.append("## References\n")
+    markdown.append(str(paper.get('references')))
+    
+    return "".join(markdown)
+
+def download_as_docx(paper: dict) -> str:
+    """Convert a paper dictionary into a DOCX file."""
+    def temp_image_file(url: str):
+        image = requests.get(url)
+        image = Image.open(BytesIO(image.content))
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.png', delete=False) as temp_file:
+            image.save(temp_file.name)
+            image_file_path = temp_file.name
+        return image_file_path
+    if paper.get('image_url'):
+        image_path = temp_image_file(paper.get('image_url'))
+    else:
+        image_path = None
+
+    # Prepare markdown
+    markdown = []
+    
+    # Add title
+    markdown.append(f"# {paper.get('title')}")
+    markdown.append("\n\n")
+    if image_path:
+        markdown.append(f"![Image]({image_path})\n")
+        markdown.append("\n\n")
+    
+    # Add body sections
+    for section in paper.get('paragraphs', []):
+        markdown.append(f"## {section.get('title')}\n")
+        if section.get('content'):
+            markdown.append(section.get('content'))
+            markdown.append("\n")
+        markdown.append("\n\n")
+    
+    if paper.get('table'):
+        try:
+            markdown.append(pd.DataFrame(data=paper.get('table')['data'], 
+                                      columns=paper.get('table')['columns']).to_markdown())
+        except:
+            markdown.append(pd.DataFrame(paper.get('table')).to_markdown())
+        markdown.append("\n\n")
+    
+    # Add references
+    markdown.append("## References\n")
+    markdown.append(str(paper.get('references')))
+    
+    markdown_str = "".join(markdown)
+    
+    # Create a temporary file
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as temp_file:
+        temp_file.write(markdown_str)
+        temp_file_path = temp_file.name
+
+    # Create a new document and load from the temporary files
+    document = Document()
+    document.LoadFromFile(temp_file_path)
+
+    # Save as DOCX
+    output_path = "output.docx"
+    document.SaveToFile(output_path, FileFormat.Docx)
+
+    # Clean up temporary files
+    os.unlink(temp_file_path)
+    if image_path:
+        os.unlink(image_path)
+    
+    return output_path
+
+async def process_query(query: str):
+    """Process a user query and update the research paper."""
+    response = await agent.chat(query)
+    st.session_state.chat_history.append({"role": "user", "content": query})
+    st.session_state.chat_history.append({"role": "assistant", "content": str(response)})
+    st.session_state.agent_memory = agent.memory.messages
+    # Update research paper if available
+    if agent.deps.deep_search_results:
+        st.session_state.research_paper = agent.deps.deep_search_results
+        st.session_state.table_data = agent.deps.table_data
+    if agent.deps.table_data:
+        st.session_state.table_data = agent.deps.table_data
+    if agent.deps.quick_search_results:
+        st.session_state.quick_search_results = agent.deps.quick_search_results
+    
+
+def process_query_sync(query: str):
+    """Synchronous wrapper for async process_query"""
+    if not st.session_state.loop.is_running():
+        return st.session_state.loop.run_until_complete(process_query(query))
+    else:
+        # Create a new event loop if the current one is running
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            return loop.run_until_complete(process_query(query))
+        finally:
+            loop.close()
+
+def reset_chat():
+    """Reset the chat history and agent state."""
+    agent.reset()
+    st.session_state.chat_history = []
+    st.session_state.research_paper = {}
+    st.session_state.table_data = {}
+    st.session_state.quick_search_results = []
+    st.session_state.agent_memory = []
+    
+
+
+# Streamlit UI
+st.title("Research Assistant Chatbot")
+
+# Create a container for the chat interface
+chat_container = st.container()
+
+# Create two columns - one for chat and one for controls
+
+
+
+# Chat interface with custom styling
+st.markdown("""
+    <style>
+    .stChatMessage {
+        padding: 20px;
+        border-radius: 10px;
+        margin-bottom: 10px;
+    }
+    .user-message {
+        background-color: #e6f3ff;
+        text-align: right;
+    }
+    .assistant-message {
+        background-color: #f0f0f0;
+        text-align: left;
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+# Chat messages in a scrollable container
+with st.container(height=300):
+    for message in st.session_state.chat_history:
+        with st.chat_message(message["role"], avatar="🧑" if message["role"] == "user" else "🤖"):
+            st.write(message["content"])
+
+# Input area at the bottom
+with st.container():
+    query = st.text_input("Ask me anything...", key="query_input", placeholder="Type your question here...")
+    col3, col4 = st.columns([6, 1])
+    with col3:
+        if st.button("Submit", use_container_width=True):
+            if query:
+                with st.spinner('Processing your request...'):
+                    process_query_sync(query)
+                    st.rerun()
+
+
+# Control buttons in the sidebar
+with st.sidebar:
+    st.markdown("### Helpful Tips")
+    st.markdown("---")
+    st.markdown("""reset chat if you want to start over or if the agent is not responding properly""")
+    if st.button("Reset Chat", use_container_width=True):
+        reset_chat()
+
+    # Add some helpful tips
+    st.markdown("---")
+    st.markdown("### Tips")
+    st.markdown("""
+    - ask the agent about its capabilities and what it can do
+    - the table is editable, you can edit in the files section or in the chat
+    - the research paper is also editable, you can edit in the chat by asking the agent to edit the research paper
+    - the agent can also add a table to the research paper, you can see the table in the files section or in the chat
+    """)
+
+st.markdown('### Files ')
+
+# Research Paper popover
+ 
+if st.session_state.research_paper :
+    with st.popover("📄 Research Paper",use_container_width=True):
+        st.markdown("""
+        <style>
+        .paper-preview {
+            border: 1px solid #e0e0e0;
+            border-radius: 5px;
+            padding: 10px;
+            margin-bottom: 10px;
+        }
+        </style>
+        """, unsafe_allow_html=True)
+        
+        # Paper content
+        st.markdown(paper_to_markdown(st.session_state.research_paper))
+        
+        # Download button
+        if st.button("📥 Save as DOCX", use_container_width=True):
+            with st.spinner('Preparing document...'):
+                output_path = download_as_docx(st.session_state.research_paper)
+                with open(output_path, "rb") as file:
+                    st.download_button(
+                        label="Download",
+                        data=file,
+                        file_name="research_paper.docx",
+                        mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+                    )
+                os.unlink(output_path)
+
+# Table popover
+if st.session_state.table_data:
+    with st.popover("📊 Editable Table",use_container_width=True):
+        try:
+            try:
+                df = pd.DataFrame(data=st.session_state.table_data['data'], 
+                                columns=st.session_state.table_data['columns'])
+            except:
+                df = pd.DataFrame(st.session_state.table_data)
+            
+
+            edited_df = st.data_editor(df)
+            if st.session_state.research_paper:
+                if st.button("💾 add table to research paper", use_container_width=True):
+                    # Convert DataFrame back to the correct structure
+                    table_dict = {
+                        'data': edited_df.values.tolist(),
+                        'columns': edited_df.columns.tolist()
+                    }
+                    st.session_state.research_paper['table'] = table_dict
+                    st.session_state.table_data = table_dict
+                    st.success("Table updated successfully!")
+                    st.rerun()
+        except Exception as e:
+            st.error(f"Error loading table data: {str(e)}")
+
+

deep_research.py

@@ -0,0 +1,185 @@
+from pydantic_graph import BaseNode, End, GraphRunContext, Graph
+from pydantic_ai import Agent
+from pydantic_ai.common_tools.tavily import tavily_search_tool
+from dataclasses import dataclass
+from pydantic import Field, BaseModel
+from typing import  List, Dict, Optional, Any
+from pydantic_ai.models.gemini import GeminiModel
+from pydantic_ai.providers.google_gla import GoogleGLAProvider
+from dotenv import load_dotenv
+import os
+from tavily import TavilyClient
+from IPython.display import Image, display
+import requests
+import time
+
+load_dotenv()
+google_api_key=os.getenv('google_api_key')
+tavily_key=os.getenv('tavily_key')
+tavily_client = TavilyClient(api_key=tavily_key)
+llm=GeminiModel('gemini-2.0-flash', provider=GoogleGLAProvider(api_key=google_api_key))
+pse=os.getenv('pse')
+
+@dataclass
+class State:
+    query:str
+    preliminary_research: str
+    research_plan: Dict
+    research_results: Dict
+    validation : str
+    final: Dict
+class paragraph_content(BaseModel):
+    title: str = Field(description='the title of the paragraph')
+    content: str = Field(description='the content of the paragraph')
+
+class paragraph(BaseModel):
+    title: str = Field(description='the title of the paragraph')
+    should_include: str = Field(description='a description of what the paragraph should include')  
+class Paper_layout(BaseModel):
+    title: str = Field(description='the title of the paper')
+    paragraphs: List[paragraph]= Field(description='the list of paragraphs of the paper')
+
+paper_layout_agent=Agent(llm, result_type=Paper_layout, system_prompt="generate a paper layout based on the query, preliminary_search, search_results,include a Title for the paper, for the paragraphs only include the title, no content, no image, no table, start with introduction and end with conclusion")
+paragraph_gen_agent=Agent(llm, result_type=paragraph_content, system_prompt="generate a paragraph synthesizing the research_results based on the title,what the paragraph should include, and what has already been written to avoid repetition")
+class PaperGen_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State])->End:
+        prompt=(f'query:{ctx.state.query}, preliminary_search:{ctx.state.preliminary_research},search_results:{ctx.state.research_results.research_results}')
+        result=await paper_layout_agent.run(prompt)
+        paragraphs=[]
+        for i in result.data.paragraphs:
+            time.sleep(2)
+            paragraph_data=await paragraph_gen_agent.run(f'title:{i.title}, should_include:{i.should_include}, research_results:{ctx.state.research_results.research_results}, already_written:{paragraphs}')
+            paragraphs.append(paragraph_data.data.model_dump())
+
+        paper={'title':result.data.title,
+                'image_url':ctx.state.research_results.image_url if ctx.state.research_results.image_url else None,
+                'paragraphs':paragraphs,
+                'table':ctx.state.research_results.table if ctx.state.research_results.table else None,
+                'references':ctx.state.research_results.references if ctx.state.research_results.references else None}
+
+        ctx.state.final=paper
+                
+        return End(ctx.state.final)
+
+def google_image_search(query:str):
+  """Search for images using Google Custom Search API
+  args: query
+  return: image url
+  """
+  # Define the API endpoint for Google Custom Search
+  url = "https://www.googleapis.com/customsearch/v1"
+
+  params = {
+      "q": query,
+      "cx": pse,
+      "key": google_api_key,
+      "searchType": "image",  # Search for images
+      "num": 1  # Number of results to fetch
+  }
+
+  # Make the request to the Google Custom Search API
+  response = requests.get(url, params=params)
+  data = response.json()
+
+  # Check if the response contains image results
+  if 'items' in data:
+      # Extract the first image result
+      image_url = data['items'][0]['link']
+      return image_url
+
+class Table_row(BaseModel):
+    data: List[str] = Field(description='the data of the row')
+class Table(BaseModel):
+    rows: List[Table_row] = Field(description='the rows of the table')
+    columns: List[str] = Field(description='the columns of the table')
+
+class Research_results(BaseModel):
+    research_results: List[str] = Field(default_factory=None,description='the research results')
+    image_url: str = Field(default_factory=None,description='the image url if needed else return None')
+    table: dict = Field(default_factory=None,description='the table dataframe in a dictionary format')
+    references: str = Field(default_factory=None,description='the references (urls) of the research_results')
+
+table_agent=Agent(llm, result_type=Table, system_prompt="generate a detailed table in dictionary format based on the research and the query")
+
+class Research_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State])->PaperGen_node:
+        research_results=Research_results(research_results=[], image_url='', table={}, references='')
+        
+        for i in ctx.state.research_plan.search_queries:
+            response = tavily_client.search(i.search_query)
+            data=[]
+            for i in response.get('results'):
+                if i.get('score')>0.50:
+                    
+                    
+                    data.append(i.get('url'))
+                    research_results.research_results.append(i.get('content'))
+        research_results.research_results=list(set(research_results.research_results))
+        research_results.references=list(set(data))
+        research_results.references=', '.join(research_results.references)
+        ctx.state.research_results=research_results
+        if ctx.state.research_plan.image_search_query:
+            image_url=google_image_search(ctx.state.research_plan.image_search_query)
+            ctx.state.research_results.image_url=image_url
+       
+        if ctx.state.research_plan.table:
+            result=await table_agent.run(f'research_results:{ctx.state.research_results.research_results},query:{ctx.state.query}')
+            ctx.state.research_results.table={'data':[row.data for row in result.data.rows], 'columns':result.data.columns}
+        
+        
+        return PaperGen_node()
+
+class search_query(BaseModel):
+    search_query: str = Field(description='the detailed web search query for the research')
+    
+class Research_plan(BaseModel):
+    search_queries: List[search_query] = Field(description='the detailed web search queries for the research')
+    table: Optional[str] = Field(default_factory=None,description='if a table is needed, return yes else return None')
+    image_search_query: Optional[str] = Field(default_factory=None,description='if image is needed, generate a image search query, optional')
+
+research_plan_agent=Agent(llm, result_type=Research_plan, system_prompt='generate a detailed research plan breaking down the research into smaller parts based on the query and the preliminary search, include a table and image search query if the user wants it')
+
+class Research_plan_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State])->Research_node:
+        
+        prompt=(f'query:{ctx.state.query}, preliminary_search:{ctx.state.preliminary_research}')
+        result=await research_plan_agent.run(prompt)
+        ctx.state.research_plan=result.data
+        return Research_node()
+    
+    
+search_agent=Agent(llm, tools=[tavily_search_tool(tavily_key)], system_prompt="do a websearch based on the query")
+
+class preliminary_search_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State]) -> Research_plan_node:
+        prompt = (' Do a preliminary search to get a global idea of the subject that the user wants to do reseach on as well as the necessary informations to do a search on.\n'
+                  f'The subject is based on the query: {ctx.state.query}, return the results of the search.')
+        result=await search_agent.run(prompt)
+        ctx.state.preliminary_research=result.data
+        return Research_plan_node()
+
+class Deep_research_engine:
+    def __init__(self):
+        self.graph=Graph(nodes=[preliminary_search_node, Research_plan_node, Research_node, PaperGen_node])
+        self.state=State(query='', preliminary_research='', research_plan=[], research_results=[], validation='', final='')
+
+    async def chat(self,query:str):
+        """Chat with the deep research engine,
+        Args:
+            query (str): The query to search for
+        Returns:
+            str: The response from the deep research engine
+        """
+        self.state.query=query
+        response=await self.graph.run(preliminary_search_node(),state=self.state)
+        return response.output
+
+
+    def display_graph(self):
+        """Display the graph of the deep research engine
+        Returns:
+            Image: The image of the graph
+        """
+        image=self.graph.mermaid_image()
+        return display(Image(image))
+

main_agent.py

@@ -0,0 +1,208 @@
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.common_tools.tavily import tavily_search_tool
+from pydantic_ai.messages import ModelMessage
+from dotenv import load_dotenv
+import os
+from pydantic import Field, BaseModel
+from typing import Dict, List, Any
+from deep_research import Deep_research_engine
+from pydantic_ai.models.gemini import GeminiModel
+from pydantic_ai.providers.google_gla import GoogleGLAProvider
+from dataclasses import dataclass
+from typing import Optional
+from spire.doc import Document,FileFormat
+from spire.doc.common import *
+import requests
+from table_maker import table_maker_engine
+from PIL import Image
+from io import BytesIO, StringIO
+import tempfile
+import pandas as pd
+
+load_dotenv()
+tavily_key=os.getenv('tavily_key')
+google_api_key=os.getenv('google_api_key')
+pse=os.getenv('pse')
+
+llm=GeminiModel('gemini-2.0-flash', provider=GoogleGLAProvider(api_key=google_api_key))
+
+
+@dataclass
+class Deps:
+    deep_search_results:dict
+    quick_search_results:list[str]
+    table_data:dict
+    
+
+
+
+async def deep_research_agent(ctx:RunContext[Deps], query:str):
+    """
+    This function is used to do a deep research on the web for information on a complex query, generates a report or a paper.
+    Args:
+        query (str): The query to search for
+    Returns:
+        str: The result of the search
+    """
+    deepsearch=Deep_research_engine()
+    res=await deepsearch.chat(query)
+    ctx.deps.deep_search_results=res
+    ctx.deps.table_data=res.get('table')
+    return str(res)
+
+quick_search_agent=Agent(llm,tools=[tavily_search_tool(tavily_key)])
+async def quick_research_agent(ctx: RunContext[Deps], query:str):
+    """
+    This function is used to do a quick search on the web for information on a given query.
+    Args:
+        query (str): The query to search for
+    Returns:
+        str: The result of the search
+    """
+    res=await quick_search_agent.run(query)
+    ctx.deps.quick_search_results.append(res.data)
+    return str(res.data)
+
+
+def google_image_search(query:str):
+  """Search for images using Google Custom Search API
+  args: query
+  return: image url
+  """
+  # Define the API endpoint for Google Custom Search
+  url = "https://www.googleapis.com/customsearch/v1"
+
+  params = {
+      "q": query,
+      "cx": pse,
+      "key": google_api_key,
+      "searchType": "image",  # Search for images
+      "num": 1  # Number of results to fetch
+  }
+
+  # Make the request to the Google Custom Search API
+  response = requests.get(url, params=params)
+  data = response.json()
+
+  # Check if the response contains image results
+  if 'items' in data:
+      # Extract the first image result
+      image_url = data['items'][0]['link']
+      return image_url
+  
+
+
+async def research_editor_tool(ctx: RunContext[Deps], query:str):
+    """
+    Use this tool to edit the deep search result to make it more accurate following the query's instructions.
+    This tool can modify paragraphs, image_url. For image_url, you need to give the query to search for the image.
+    Args:
+        query (str): The query containing instructions for editing the deep search result
+    Returns:
+        str: The edited and improved deep search result
+    """
+    @dataclass
+    class edit_route:
+        paragraph_number:Optional[int] = Field(default_factory=None, description='the number of the paragraph to edit, if the paragraph is not needed to be edited, return None')
+        route: str = Field(description='the route to the content to edit, either paragraphs, image_url')
+
+ 
+   
+    paper_dict={'title':ctx.deps.deep_search_results.get('title'),
+                                'image_url':ctx.deps.deep_search_results.get('image_url') if ctx.deps.deep_search_results.get('image_url') else 'None',
+                                'paragraphs_title':{num:paragraph.get('title') for num,paragraph in enumerate(ctx.deps.deep_search_results.get('paragraphs'))}, 
+                                'table':ctx.deps.deep_search_results.get('table') if ctx.deps.deep_search_results.get('table') else 'None',
+                                'references':ctx.deps.deep_search_results.get('references')}
+    
+    route_agent=Agent(llm,result_type=edit_route, system_prompt="you decide the route to the content to edit based on the query's instructions and the paper_dict, either paragraphs, image_url")
+    route=await route_agent.run(f'query:{query}, paper_dict:{paper_dict}')
+    contents=ctx.deps.deep_search_results
+ 
+  
+    @dataclass
+    class Research_edits:
+        edits:str = Field(description='the edits')
+    editor_agent=Agent(llm,tools=[google_image_search],result_type=Research_edits, system_prompt="you are an editor, you are given a query, some content to edit, and maybe a quick search result (optional), you need to edit the deep search result to make it more accurate following the query's instructions, return only the edited content, no comments")
+    if route.data.route=='paragraphs':
+        content=contents.get('paragraphs')[route.data.paragraph_number]['content']
+        res=await editor_agent.run(f'query:{query}, content:{content}, quick_search_results:{ctx.deps.quick_search_results if ctx.deps.quick_search_results else "None"}')
+        ctx.deps.deep_search_results['paragraphs'][route.data.paragraph_number]['content']=res.data.edits
+    if route.data.route=='image_url':
+        content=contents.get('image_url')
+        res=await editor_agent.run(f'query:{query}, content:{content}, quick_search_results:{ctx.deps.quick_search_results if ctx.deps.quick_search_results else "None"}')
+        ctx.deps.deep_search_results['image_url']=res.data.edits
+
+    
+    return str(ctx.deps.deep_search_results)
+
+
+async def Table_agent(ctx: RunContext[Deps], query:str):
+    """
+    Use this tool to create a table, edit a table or add a table to the deep search result. the add table to paper route is used to create and add a table to the deep search result.
+    Args:
+        query (str): The query to create a table, edit a table or add a table to the deep search result
+    Returns:
+        dict: The table
+    """
+    @dataclass
+    class route:
+        route: str = Field(description='the route to the content to edit, either create_table, edit_table, or add_table_to_paper')
+    route_agent=Agent(llm,result_type=route, system_prompt="you decide the route to the content to edit based on the query's instructions, return only the route, either create_table, edit_table, or add_table_to_paper")
+    route=await route_agent.run(f'query:{query}')
+
+
+    if route.data.route=='create_table':
+        table_maker=table_maker_engine()
+        table=await table_maker.chat(query)
+        ctx.deps.table_data=table
+        return str(table)
+    
+    if route.data.route=='edit_table':
+        table=ctx.deps.table_data
+        class Table_row(BaseModel):
+            data: List[str] = Field(description='the data of the row')
+        class Table(BaseModel): 
+            rows: List[Table_row] = Field(description='the rows of the table')
+            columns: List[str] = Field(description='the columns of the table')
+
+        table_editor=Agent(llm, result_type=Table, system_prompt="edit the table based on the query's instructions, the research results (if any) and the quick search results(if any)")
+        generated_table=await table_editor.run(f'query:{query}, table:{table}, research:{ctx.deps.deep_search_results if ctx.deps.deep_search_results else "None"}, quick_search_results:{ctx.deps.quick_search_results if ctx.deps.quick_search_results else "None"}')
+        ctx.deps.table_data={'data':[row.data for row in generated_table.data.rows], 'columns':generated_table.data.columns}
+        return str(ctx.deps.table_data)
+    
+    if route.data.route=='add_table_to_paper':
+        class Table_row(BaseModel):
+            data: List[str] = Field(description='the data of the row')
+        class Table(BaseModel): 
+            rows: List[Table_row] = Field(description='the rows of the table')
+            columns: List[str] = Field(description='the columns of the table')
+        table_creator=Agent(llm, result_type=Table, system_prompt="create a table based on the query's instructions, the research results (if any) and the quick search results(if any)")
+        generated_table=await table_creator.run(f'query:{query}, research:{ctx.deps.deep_search_results if ctx.deps.deep_search_results else "None"}, quick_search_results:{ctx.deps.quick_search_results if ctx.deps.quick_search_results else "None"}')
+        ctx.deps.deep_search_results['table']={'data':[row.data for row in generated_table.data.rows], 'columns':generated_table.data.columns}
+        ctx.deps.table_data=ctx.deps.deep_search_results['table']
+        return str(ctx.deps.deep_search_results)
+
+@dataclass
+class Message_state:
+    messages: list[ModelMessage]
+
+
+
+class Main_agent:
+    def __init__(self):
+        self.agent=Agent(llm, system_prompt="you are a research assistant, you are given a query, leverage what tool(s) to use, make suggestions to the user about the tools to use, \
+                          never show the output of the tools, except for the table, notify the user about what next step they can take, inform the user about the table,\
+                         and the table's editable nature either in the chat or in the files section",
+                          tools=[deep_research_agent,research_editor_tool,quick_research_agent,Table_agent])
+        self.deps=Deps( deep_search_results=[], quick_search_results=[], table_data={})
+        self.memory=Message_state(messages=[])
+
+    async def chat(self, query:str):
+        result = await self.agent.run(query,deps=self.deps, message_history=self.memory.messages)
+        self.memory.messages=result.all_messages()
+        return result.data
+    
+    def reset(self):
+        self.memory.messages=[]
+        self.deps=Deps( deep_search_results=[], quick_search_results=[])
+

requirements.txt

@@ -0,0 +1,14 @@
+ipython==9.0.2
+Pillow==11.1.0
+pydantic==2.11.2
+pydantic_ai==0.0.52
+pydantic_graph==0.0.52
+python-dotenv==1.1.0
+Requests==2.32.3
+Spire.Doc.Free==12.12.0
+tavily_python==0.5.1
+tabulate==0.9.0
+pandas==2.1.1
+python-pptx
+streamlit==1.26.0
+nest_asyncio==1.6.0

table_maker.py

@@ -0,0 +1,106 @@
+from pydantic_graph import BaseNode, End, GraphRunContext, Graph
+from pydantic_ai import Agent
+from pydantic_ai.common_tools.tavily import tavily_search_tool
+from dataclasses import dataclass
+from pydantic import Field, BaseModel
+from typing import  List, Dict, Optional, Any
+from pydantic_ai.models.gemini import GeminiModel
+from pydantic_ai.providers.google_gla import GoogleGLAProvider
+from dotenv import load_dotenv
+import os
+from tavily import TavilyClient
+from IPython.display import Image, display
+import requests
+
+load_dotenv()
+google_api_key=os.getenv('google_api_key')
+tavily_key=os.getenv('tavily_key')
+tavily_client = TavilyClient(api_key=tavily_key)
+llm=GeminiModel('gemini-2.0-flash', provider=GoogleGLAProvider(api_key=google_api_key))
+
+@dataclass
+class State:
+    query:str
+    research:List[str]
+    table:dict
+    preliminary_research:str
+    research_plan:List[str]
+
+#define the table row and table schema
+class Table_row(BaseModel):
+    data: List[str] = Field(description='the data of the row')
+class Table(BaseModel): 
+    rows: List[Table_row] = Field(description='the rows of the table')
+    columns: List[str] = Field(description='the columns of the table')
+    
+
+    
+class table_maker_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State])->End:
+        table_agent=Agent(llm, result_type=Table, system_prompt="generate a detailed table in a dictionary format based on the research and the query")
+        table=await table_agent.run(f'query:{ctx.state.query}, research:{ctx.state.research}')
+        ctx.state.table={'data':[row.data for row in table.data.rows], 'columns':table.data.columns}
+        return End(ctx.state.table)
+    
+
+class data_research_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State])->table_maker_node:
+        for i in ctx.state.research_plan:   
+            response = tavily_client.search(i.search_query)
+
+            for i in response.get('results'):
+                if i.get('score')>0.50:
+                    ctx.state.research.append(i.get('content'))
+        return table_maker_node()
+    
+
+
+class search_query(BaseModel):
+    search_query: str = Field(description='the detailed web search query for the research')
+    
+class Research_plan(BaseModel):
+    search_queries: List[search_query] = Field(description='the detailed web search queries for the research')
+
+research_plan_agent=Agent(llm, result_type=Research_plan, system_prompt='generate a detailed research plan breaking down the research into smaller parts based on the query and the preliminary search')
+
+class Research_plan_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State])->data_research_node:
+        
+        prompt=(f'query:{ctx.state.query}, preliminary_search:{ctx.state.preliminary_research}')
+        result=await research_plan_agent.run(prompt)
+        ctx.state.research_plan=result.data.search_queries
+        return data_research_node()
+
+class preliminary_search_node(BaseNode[State]):
+    async def run(self, ctx: GraphRunContext[State]) -> Research_plan_node:
+        search_agent=Agent(llm, tools=[tavily_search_tool(tavily_key)], system_prompt="do a websearch based on the query")
+        prompt = (' Do a preliminary search to get a global idea of the subject that the user wants to do reseach on as well as the necessary informations to do a search on.\n'
+                  f'The subject is based on the query: {ctx.state.query}, return the results of the search.')
+        result=await search_agent.run(prompt)
+        ctx.state.preliminary_research=result.data
+        return Research_plan_node()
+
+    
+class table_maker_engine:
+    def __init__(self):
+        self.graph=Graph(nodes=[preliminary_search_node, Research_plan_node, data_research_node, table_maker_node])
+        self.state=State(query='', research=[], table={}, preliminary_research='', research_plan=[])
+
+    async def chat(self,query:str):
+        """Chat with the table maker engine,
+        Args:
+            query (str): The query to search for
+        Returns:
+            str: The response from the table maker engine
+        """
+        self.state.query=query
+        response=await self.graph.run(preliminary_search_node(),state=self.state)
+        return response.output
+    
+    def display_graph(self):
+        """Display the graph of the table maker engine
+        Returns:
+            Image: The image of the graph
+        """
+        image=self.graph.mermaid_image()
+        return display(Image(image))