aiagents4pharma 1.20.1__py3-none-any.whl → 1.22.0__py3-none-any.whl

aiagents4pharma/talk2scholars/agents/s2_agent.py

@@ -5,37 +5,80 @@ Agent for interacting with Semantic Scholar
 """
 
 import logging
+from typing import Any, Dict
 import hydra
 from langchain_openai import ChatOpenAI
+from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import create_react_agent, ToolNode
 from langgraph.checkpoint.memory import MemorySaver
 from ..state.state_talk2scholars import Talk2Scholars
 from ..tools.s2.search import search_tool as s2_search
 from ..tools.s2.display_results import display_results as s2_display
+from ..tools.s2.query_results import query_results as s2_query_results
+from ..tools.s2.retrieve_semantic_scholar_paper_id import (
+    retrieve_semantic_scholar_paper_id as s2_retrieve_id,
+)
 from ..tools.s2.single_paper_rec import (
     get_single_paper_recommendations as s2_single_rec,
 )
 from ..tools.s2.multi_paper_rec import get_multi_paper_recommendations as s2_multi_rec
 
-
 # Initialize logger
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def get_app(uniq_id, llm_model="gpt-4o-mini"):
+def get_app(
+    uniq_id, llm_model: BaseChatModel = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+):
     """
-    This function returns the langraph app.
+    Initializes and returns the LangGraph application for the Semantic Scholar (S2) agent.
+
+    This function sets up the S2 agent, which integrates various tools to search, retrieve,
+    and display research papers from Semantic Scholar. The agent follows the ReAct pattern
+    for structured interaction.
+
+    Args:
+        uniq_id (str): Unique identifier for the current conversation session.
+        llm_model (BaseChatModel, optional): The language model to be used by the agent.
+            Defaults to `ChatOpenAI(model="gpt-4o-mini", temperature=0)`.
+
+    Returns:
+        StateGraph: A compiled LangGraph application that enables the S2 agent to process
+            user queries and retrieve research papers.
+
+    Example:
+        >>> app = get_app("thread_123")
+        >>> result = app.invoke(initial_state)
     """
 
-    def agent_s2_node(state: Talk2Scholars):
+    # def agent_s2_node(state: Talk2Scholars) -> Command[Literal["supervisor"]]:
+    def agent_s2_node(state: Talk2Scholars) -> Dict[str, Any]:
         """
-        This function calls the model.
+        Processes the user query and retrieves relevant research papers.
+
+        This function calls the language model using the configured `ReAct` agent to analyze
+        the state and generate an appropriate response. The function then returns control
+        to the main supervisor.
+
+        Args:
+            state (Talk2Scholars): The current conversation state, including messages exchanged
+                and any previously retrieved research papers.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the updated conversation state.
+
+        Example:
+            >>> result = agent_s2_node(current_state)
+            >>> papers = result.get("papers", [])
         """
         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
+        result = model.invoke(state, {"configurable": {"thread_id": uniq_id}})
+
+        return result
+
+    logger.log(logging.INFO, "thread_id, llm_model: %s, %s", uniq_id, llm_model)
 
     # Load hydra configuration
     logger.log(logging.INFO, "Load Hydra configuration for Talk2Scholars S2 agent.")
@@ -46,30 +89,31 @@ def get_app(uniq_id, llm_model="gpt-4o-mini"):
         cfg = cfg.agents.talk2scholars.s2_agent
 
     # Define the tools
-    tools = ToolNode([s2_search, s2_display, s2_single_rec, s2_multi_rec])
+    tools = ToolNode(
+        [
+            s2_search,
+            s2_display,
+            s2_query_results,
+            s2_retrieve_id,
+            s2_single_rec,
+            s2_multi_rec,
+        ]
+    )
 
     # Define the model
     logger.log(logging.INFO, "Using OpenAI model %s", llm_model)
-    llm = ChatOpenAI(model=llm_model, temperature=cfg.temperature)
 
     # Create the agent
     model = create_react_agent(
-        llm,
+        llm_model,
         tools=tools,
         state_schema=Talk2Scholars,
-        # prompt=cfg.s2_agent,
         state_modifier=cfg.s2_agent,
         checkpointer=MemorySaver(),
     )
 
-    # Define a new graph
     workflow = StateGraph(Talk2Scholars)
-
-    # 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

aiagents4pharma/talk2scholars/configs/agents/talk2scholars/main_agent/default.yaml

@@ -5,14 +5,35 @@ openai_llms:
   - "gpt-4-turbo"
   - "gpt-3.5-turbo"
 temperature: 0
-main_agent: >
-  You are an intelligent research assistant coordinating academic paper discovery and analysis.
+system_prompt: >
+  You are the Talk2Scholars agent coordinating academic paper discovery and analysis.
 
-  AVAILABLE TOOLS AND ROUTING:
-  1. semantic_scholar_agent:
-     Access to tools:
-     - search_tool: For paper discovery
-     - display_results: For showing paper results
-     - get_single_paper_recommendations: For single paper recommendations
-     - get_multi_paper_recommendations: For multi-paper recommendations
-     → ROUTE TO THIS AGENT FOR: Any query about academic papers, research, or articles
+  You have access to the following agents:
+  1. S2_agent: This agent can be used to search and recommend papers
+    from Semantic Scholar. Use this agent when the user asks for
+    general paper searches and recommendations. This agent can also
+    retrieve the Semantic Scholar ID of a paper.
+router_prompt: >
+  You are a supervisor tasked with managing a conversation between the
+  following workers: {members}. Given the user request, respond with the
+  worker to act next. Each worker will perform a task and respond with
+  their results and status. When finished, respond with FINISH.
+
+  Here is a description of the workers:
+  1. S2_agent: This agent can be used to search and recommend papers
+    from Semantic Scholar. Use this agent when the user asks for
+    general paper searches and recommendations. This agent can also
+    retrieve the Semantic Scholar ID of a paper. It can also be used to
+    provide more information about a paper.
+
+  Here are some instructions for the workers:
+  1. Call the S2 agent for general paper searches and recommendations.
+  2. The S2 agent has access to tools for querying and displaying papers.
+  3. If the user wants suggestions for papers and you don’t have 
+    a Semantic Scholar ID for it but do have the title from
+    the last displayed results, use the S2 agent to retrieve the
+    Semantic Scholar ID of the paper. Then, use the S2 agent again to display
+    recommendations for the paper.
+  4. You can call the S2 agent to get more information about a paper based
+    on the context of the conversation.
+  5. Respond with FINISH when all tasks are completed.

aiagents4pharma/talk2scholars/configs/agents/talk2scholars/s2_agent/default.yaml

@@ -6,19 +6,11 @@ openai_llms:
   - "gpt-3.5-turbo"
 temperature: 0
 s2_agent: >
-  You are a specialized academic research agent with access to tools for paper discovery and analysis.
-
-  YOUR TOOLS:
-  1. search_tool:
-     - Finds research papers based on user queries.
-     - If no papers are found, it performs a new search.
-
-  2. display_results:
-     - Shows the current research papers.
-     - If no papers are found, it will instruct you to perform a search.
-
-  3. get_single_paper_recommendations:
-     - Provides recommendations based on a single selected paper.
-
-  4. get_multi_paper_recommendations:
-     - Provides recommendations based on multiple selected papers.
+  You are an academic research assistant with access to the
+  Semantic Scholar API for paper discovery and analysis.
+  You also have tools to gain more insights on the papers and
+  display them.
+  You must strictly rely on retrieved information and avoid
+  generating unsupported content. Do not generate hallucinations
+  or fabricate details of any article. Stay focused on accurate,
+  sourced academic insights.

aiagents4pharma/talk2scholars/configs/app/frontend/default.yaml

@@ -1,14 +1,13 @@
-# # Page configuration
-# page:
-#   title: "Talk2Scholars"
-#   icon: "🤖"
-#   layout: "wide"
+# Page configuration
+page:
+  title: "Talk2Scholars"
+  icon: "🤖"
+  layout: "wide"
 
 # Available LLM models
-llm_models:
-  - "gpt-4o-mini"
-  - "gpt-4-turbo"
-  - "gpt-3.5-turbo"
+llms:
+  available_models:
+    - "OpenAI/gpt-4o-mini"
 # # Chat UI configuration
 # chat:
 #   assistant_avatar: "🤖"
@@ -16,6 +15,9 @@ llm_models:
 #   input_placeholder: "Say something ..."
 #   spinner_text: "Fetching response ..."
 
+api_keys:
+  openai_key: "OPENAI_API_KEY"
+  nvidia_key: "NVIDIA_API_KEY"
 # # Feedback configuration
 # feedback:
 #   type: "thumbs"

aiagents4pharma/talk2scholars/configs/config.yaml

@@ -5,4 +5,5 @@ defaults:
   - tools/search: default
   - tools/single_paper_recommendation: default
   - tools/multi_paper_recommendation: default
+  - tools/retrieve_semantic_scholar_paper_id: default
   - app/frontend: default

aiagents4pharma/talk2scholars/configs/tools/multi_paper_recommendation/default.yaml

@@ -9,6 +9,8 @@ api_fields:
   - "authors"
   - "citationCount"
   - "url"
+# Commented fields that could be added later if needed
+# - "externalIds"
 
 # Default headers and params
 headers:

aiagents4pharma/talk2scholars/configs/tools/retrieve_semantic_scholar_paper_id/__init__.py

@@ -0,0 +1,3 @@
+"""
+Import all the modules in the package
+"""

aiagents4pharma/talk2scholars/configs/tools/search/default.yaml

@@ -10,6 +10,7 @@ api_fields:
   - "citationCount"
   - "url"
 # Commented fields that could be added later if needed
+# - "externalIds"
 #  - "publicationTypes"
 #  - "openAccessPdf"
 

aiagents4pharma/talk2scholars/configs/tools/single_paper_recommendation/default.yaml

@@ -10,6 +10,7 @@ api_fields:
   - "citationCount"
   - "url"
 # Commented fields that could be added later if needed
+# - "externalIds"
 #  - "publicationTypes"
 #  - "openAccessPdf"
 

aiagents4pharma/talk2scholars/state/state_talk2scholars.py

@@ -1,9 +1,14 @@
 """
-This is the state file for the talk2scholars agent.
+State management for the Talk2Scholars agent.
+
+This module defines the state class `Talk2Scholars`, which maintains the conversation
+context, retrieved papers, and other relevant metadata. The state ensures consistency
+across agent interactions.
 """
 
 import logging
 from typing import Annotated, Any, Dict
+from langchain_core.language_models import BaseChatModel
 from langgraph.prebuilt.chat_agent_executor import AgentState
 
 # Configure logging
@@ -12,22 +17,46 @@ logger = logging.getLogger(__name__)
 
 
 def replace_dict(existing: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, Any]:
-    """Replace the existing dict with the new one."""
+    """
+    Replaces the existing dictionary with a new dictionary.
+
+    This function logs the state update and ensures that the old state is replaced
+    with the new one.
+
+    Args:
+        existing (Dict[str, Any]): The current dictionary state.
+        new (Dict[str, Any]): The new dictionary state to replace the existing one.
+
+    Returns:
+        Dict[str, Any]: The updated dictionary state.
+
+    Example:
+        >>> old_state = {"papers": {"id1": "Paper 1"}}
+        >>> new_state = {"papers": {"id2": "Paper 2"}}
+        >>> updated_state = replace_dict(old_state, new_state)
+        >>> print(updated_state)
+        {"papers": {"id2": "Paper 2"}}
+    """
     logger.info("Updating existing state %s with the state dict: %s", existing, new)
     return new
 
 
 class Talk2Scholars(AgentState):
     """
-    The state for the talk2scholars agent, inheriting from AgentState.
+    Represents the state of the Talk2Scholars agent.
+
+    This class extends `AgentState` to maintain conversation history, retrieved papers,
+    and interactions with the language model.
 
     Attributes:
-        papers: Dictionary of papers from search results
-        multi_papers: Dictionary of papers from multi-paper recommendations
-        llm_model: Model being used
+        last_displayed_papers (Dict[str, Any]): Stores the most recently displayed papers.
+        papers (Dict[str, Any]): Stores the research papers retrieved from the agent's queries.
+        multi_papers (Dict[str, Any]): Stores multiple recommended papers from various sources.
+        llm_model (BaseChatModel): The language model instance used for generating responses.
     """
 
     # Agent state fields
+    last_displayed_papers: Annotated[Dict[str, Any], replace_dict]
     papers: Annotated[Dict[str, Any], replace_dict]
     multi_papers: Annotated[Dict[str, Any], replace_dict]
-    llm_model: str
+    llm_model: BaseChatModel

aiagents4pharma/talk2scholars/tests/test_llm_main_integration.py

@@ -0,0 +1,58 @@
+"""
+Integration tests for talk2scholars system with OpenAI.
+"""
+
+import os
+import pytest
+import hydra
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, AIMessage
+from ..agents.main_agent import get_app
+from ..state.state_talk2scholars import Talk2Scholars
+
+# pylint: disable=redefined-outer-name
+
+
+@pytest.mark.skipif(
+    not os.getenv("OPENAI_API_KEY"), reason="Requires OpenAI API key to run"
+)
+def test_main_agent_real_llm():
+    """
+    Test that the main agent invokes S2 agent correctly
+    and updates the state with real LLM execution.
+    """
+
+    # Load Hydra Configuration EXACTLY like in main_agent.py
+    with hydra.initialize(version_base=None, config_path="../configs"):
+        cfg = hydra.compose(
+            config_name="config", overrides=["agents/talk2scholars/main_agent=default"]
+        )
+    hydra_cfg = cfg.agents.talk2scholars.main_agent
+
+    assert hydra_cfg is not None, "Hydra config failed to load"
+
+    # Use the real OpenAI API (ensure env variable is set)
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=hydra_cfg.temperature)
+
+    # Initialize main agent workflow (WITH real Hydra config)
+    thread_id = "test_thread"
+    app = get_app(thread_id, llm)
+
+    # Provide an actual user query
+    initial_state = Talk2Scholars(
+        messages=[HumanMessage(content="Find AI papers on transformers")]
+    )
+
+    # Invoke the agent (triggers supervisor → s2_agent)
+    result = app.invoke(
+        initial_state,
+        {"configurable": {"config_id": thread_id, "thread_id": thread_id}},
+    )
+
+    # Assert that the supervisor routed correctly
+    assert "messages" in result, "Expected messages in response"
+
+    # Fix: Accept AIMessage as a valid response type
+    assert isinstance(
+        result["messages"][-1], (HumanMessage, AIMessage, str)
+    ), "Last message should be a valid response"