Feat/setup talk2comp

aiagents4pharma/talk2competitors/agents/__init__.py

@@ -0,0 +1,6 @@
+'''
+This file is used to import all the modules in the package.
+'''
+
+from . import main_agent
+from . import s2_agent

aiagents4pharma/talk2competitors/agents/main_agent.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+
+"""
+Main agent for the talk2competitors app.
+"""
+
+import logging
+from typing import Literal
+from dotenv import load_dotenv
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import END, START, StateGraph
+from langgraph.types import Command
+from ..agents import s2_agent
+from ..config.config import config
+from ..state.state_talk2competitors import Talk2Competitors
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+def make_supervisor_node(llm: BaseChatModel) -> str:
+    """
+    Creates a supervisor node following LangGraph patterns.
+
+    Args:
+        llm (BaseChatModel): The language model to use for generating responses.
+
+    Returns:
+        str: The supervisor node function.
+    """
+    # options = ["FINISH", "s2_agent"]
+
+    def supervisor_node(state: Talk2Competitors) -> Command[Literal["s2_agent", "__end__"]]:
+        """
+        Supervisor node that routes to appropriate sub-agents.
+
+        Args:
+            state (Talk2Competitors): The current state of the conversation.
+
+        Returns:
+            Command[Literal["s2_agent", "__end__"]]: The command to execute next.
+        """
+        logger.info("Supervisor node called")
+
+        messages = [{"role": "system", "content": config.MAIN_AGENT_PROMPT}] + state[
+            "messages"
+        ]
+        response = llm.invoke(messages)
+        goto = (
+            "FINISH"
+            if not any(
+                kw in state["messages"][-1].content.lower()
+                for kw in ["search", "paper", "find"]
+            )
+            else "s2_agent"
+        )
+
+        if goto == "FINISH":
+            return Command(
+                goto=END,
+                update={
+                    "messages": state["messages"]
+                    + [AIMessage(content=response.content)],
+                    "is_last_step": True,
+                    "current_agent": None,
+                },
+            )
+
+        return Command(
+            goto="s2_agent",
+            update={
+                "messages": state["messages"],
+                "is_last_step": False,
+                "current_agent": "s2_agent",
+            },
+        )
+
+    return supervisor_node
+
+def get_app(thread_id: str, llm_model ='gpt-4o-mini') -> StateGraph:
+    """
+    Returns the langraph app with hierarchical structure.
+
+    Args:
+        thread_id (str): The thread ID for the conversation.
+
+    Returns:
+        The compiled langraph app.
+    """
+    def call_s2_agent(state: Talk2Competitors) -> Command[Literal["__end__"]]:
+        """
+        Node for calling the S2 agent.
+
+        Args:
+            state (Talk2Competitors): The current state of the conversation.
+
+        Returns:
+            Command[Literal["__end__"]]: The command to execute next.
+        """
+        logger.info("Calling S2 agent")
+        app = s2_agent.get_app(thread_id, llm_model)
+        response = app.invoke(state)
+        logger.info("S2 agent completed")
+        return Command(
+            goto=END,
+            update={
+                "messages": response["messages"],
+                "papers": response.get("papers", []),
+                "is_last_step": True,
+                "current_agent": "s2_agent",
+            },
+        )
+    llm = ChatOpenAI(model=llm_model, temperature=0)
+    workflow = StateGraph(Talk2Competitors)
+
+    supervisor = make_supervisor_node(llm)
+    workflow.add_node("supervisor", supervisor)
+    workflow.add_node("s2_agent", call_s2_agent)
+
+    # Define edges
+    workflow.add_edge(START, "supervisor")
+    workflow.add_edge("s2_agent", END)
+
+    app = workflow.compile(checkpointer=MemorySaver())
+    logger.info("Main agent workflow compiled")
+    return app

aiagents4pharma/talk2competitors/agents/s2_agent.py

@@ -0,0 +1,75 @@
+#/usr/bin/env python3
+
+'''
+Agent for interacting with Semantic Scholar
+'''
+
+import logging
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import create_react_agent
+from langgraph.checkpoint.memory import MemorySaver
+from ..config.config import config
+from ..state.state_talk2competitors import Talk2Competitors
+# from ..tools.s2 import s2_tools
+from ..tools.s2.search import search_tool
+from ..tools.s2.display_results import display_results
+from ..tools.s2.single_paper_rec import get_single_paper_recommendations
+from ..tools.s2.multi_paper_rec import get_multi_paper_recommendations
+
+load_dotenv()
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def get_app(uniq_id, llm_model='gpt-4o-mini'):
+    '''
+    This function returns the langraph app.
+    '''
+    def agent_s2_node(state: Talk2Competitors):
+        '''
+        This function calls the model.
+        '''
+        logger.log(logging.INFO, "Creating Agent_S2 node with thread_id %s", uniq_id)
+        response = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+        return response
+
+    # Define the tools
+    tools = [search_tool,
+            display_results,
+            get_single_paper_recommendations,
+            get_multi_paper_recommendations]
+
+    # Create the LLM
+    llm = ChatOpenAI(model=llm_model, temperature=0)
+    model = create_react_agent(
+                            llm,
+                            tools=tools,
+                            state_schema=Talk2Competitors,
+                            state_modifier=config.S2_AGENT_PROMPT,
+                            checkpointer=MemorySaver()
+                        )
+
+    # Define a new graph
+    workflow = StateGraph(Talk2Competitors)
+
+    # Define the two nodes we will cycle between
+    workflow.add_node("agent_s2", agent_s2_node)
+
+    # Set the entrypoint as `agent`
+    # This means that this node is the first one called
+    workflow.add_edge(START, "agent_s2")
+
+    # Initialize memory to persist state between graph runs
+    checkpointer = MemorySaver()
+
+    # Finally, we compile it!
+    # This compiles it into a LangChain Runnable,
+    # meaning you can use it as you would any other runnable.
+    # Note that we're (optionally) passing the memory when compiling the graph
+    app = workflow.compile(checkpointer=checkpointer)
+    logger.log(logging.INFO, "Compiled the graph")
+
+    return app

aiagents4pharma/talk2competitors/config/config.py

@@ -0,0 +1,142 @@
+class Config:
+    # API Endpoints
+    SEMANTIC_SCHOLAR_API = "https://api.semanticscholar.org/graph/v1"
+
+    # API Keys
+    SEMANTIC_SCHOLAR_API_KEY = "YOUR_API_KEY"  # Get this from Semantic Scholar
+
+    # Updated System Prompts
+    MAIN_AGENT_PROMPT = """You are a supervisory AI agent that routes user queries to specialized tools.
+Your task is to select the most appropriate tool based on the user's request.
+
+Available tools and their capabilities:
+
+1. semantic_scholar_agent:
+   - Search for academic papers and research
+   - Get paper recommendations
+   - Find similar papers
+   USE FOR: Any queries about finding papers, academic research, or getting paper recommendations
+
+2. zotero_agent:
+   - Manage paper library
+   - Save and organize papers
+   USE FOR: Saving papers or managing research library
+
+3. pdf_agent:
+   - Analyze PDF content
+   - Answer questions about documents
+   USE FOR: Analyzing or asking questions about PDF content
+
+4. arxiv_agent:
+   - Download papers from arXiv
+   USE FOR: Specifically downloading papers from arXiv
+
+ROUTING GUIDELINES:
+
+ALWAYS route to semantic_scholar_agent for:
+- Finding academic papers
+- Searching research topics
+- Getting paper recommendations
+- Finding similar papers
+- Any query about academic literature
+
+Route to zotero_agent for:
+- Saving papers to library
+- Managing references
+- Organizing research materials
+
+Route to pdf_agent for:
+- PDF content analysis
+- Document-specific questions
+- Understanding paper contents
+
+Route to arxiv_agent for:
+- Paper download requests
+- arXiv-specific queries
+
+Approach:
+1. Identify the core need in the user's query
+2. Select the most appropriate tool based on the guidelines above
+3. If unclear, ask for clarification
+4. For multi-step tasks, focus on the immediate next step
+
+Remember:
+- Be decisive in your tool selection
+- Focus on the immediate task
+- Default to semantic_scholar_agent for any paper-finding tasks
+- Ask for clarification if the request is ambiguous
+
+IMPORTANT GUIDELINES FOR PAPER RECOMMENDATIONS:
+
+For Multiple Papers:
+- When getting recommendations for multiple papers, always use get_multi_paper_recommendations tool
+- DO NOT call get_single_paper_recommendations multiple times
+- Always pass all paper IDs in a single call to get_multi_paper_recommendations
+- Use for queries like "find papers related to both/all papers" or "find similar papers to these papers"
+
+For Single Paper:
+- Use get_single_paper_recommendations when focusing on one specific paper
+- Pass only one paper ID at a time
+- Use for queries like "find papers similar to this paper" or "get recommendations for paper X"
+- Do not use for multiple papers
+
+Examples:
+- For "find related papers for both papers":
+  ✓ Use get_multi_paper_recommendations with both paper IDs
+  × Don't make multiple calls to get_single_paper_recommendations
+
+- For "find papers related to the first paper":
+  ✓ Use get_single_paper_recommendations with just that paper's ID
+  × Don't use get_multi_paper_recommendations
+
+Remember:
+- Be precise in identifying which paper ID to use for single recommendations
+- Don't reuse previous paper IDs unless specifically requested
+- For fresh paper recommendations, always use the original paper ID"""
+
+    S2_AGENT_PROMPT = """You are a specialized academic research assistant with access to the following tools:
+
+1. search_papers: 
+   USE FOR: General paper searches
+   - Enhances search terms automatically
+   - Adds relevant academic keywords
+   - Focuses on recent research when appropriate
+
+2. get_single_paper_recommendations:
+   USE FOR: Finding papers similar to a specific paper
+   - Takes a single paper ID
+   - Returns related papers
+
+3. get_multi_paper_recommendations:
+   USE FOR: Finding papers similar to multiple papers
+   - Takes multiple paper IDs
+   - Finds papers related to all inputs
+
+GUIDELINES:
+
+For paper searches:
+- Enhance search terms with academic language
+- Include field-specific terminology
+- Add "recent" or "latest" when appropriate
+- Keep queries focused and relevant
+
+For paper recommendations:
+- Identify paper IDs (40-character hexadecimal strings)
+- Use single_paper_recommendations for one ID
+- Use multi_paper_recommendations for multiple IDs
+
+Best practices:
+1. Start with a broad search if no paper IDs are provided
+2. Look for paper IDs in user input
+3. Enhance search terms for better results
+4. Consider the academic context
+5. Be prepared to refine searches based on feedback
+
+Remember:
+- Always select the most appropriate tool
+- Enhance search queries naturally
+- Consider academic context
+- Focus on delivering relevant results"""
+
+
+config = Config()

aiagents4pharma/talk2competitors/state/__init__.py

@@ -0,0 +1,5 @@
+'''
+This file is used to import all the modules in the package.
+'''
+
+from . import state_talk2competitors

aiagents4pharma/talk2competitors/state/state_talk2competitors.py

@@ -0,0 +1,31 @@
+"""
+This is the state file for the talk2comp agent.
+"""
+
+import logging
+from typing import Annotated, List, Optional
+from langgraph.prebuilt.chat_agent_executor import AgentState
+from typing_extensions import NotRequired, Required
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def replace_list(existing: List[str], new: List[str]) -> List[str]:
+    """Replace the existing list with the new one."""
+    logger.info("Updating existing state %s with the state list: %s",
+                existing,
+                new)
+    return new
+
+
+class Talk2Competitors(AgentState):
+    """
+    The state for the talk2comp agent, inheriting from AgentState.
+    """
+    papers: Annotated[List[str], replace_list]
+    search_table: NotRequired[str]
+    next: str  # Required for routing in LangGraph
+    current_agent: NotRequired[Optional[str]]
+    is_last_step: Required[bool]  # Required field for LangGraph
+    llm_model: str