bioguider 0.2.19__py3-none-any.whl → 0.2.20__py3-none-any.whl

bioguider/agents/agent_utils.py

@@ -289,8 +289,9 @@ class CustomOutputParser(AgentOutputParser):
         action_input = match.group(2)
         # Return the action and action input
         action_dict = None
-        action_input = action_input.strip(" ").strip('"')
-        action_input_replaced = action_input.replace("'", '"')
+        action_input_replaced = action_input.strip().strip(" ").strip('"').strip('`').strip()
+        action_input_replaced = action_input_replaced.replace("'", '"')
+        action_input_replaced = action_input_replaced.replace("`", '"')
         try:
             action_dict = json.loads(action_input_replaced)
         except json.JSONDecodeError:
@@ -299,10 +300,11 @@ class CustomOutputParser(AgentOutputParser):
             # try using ast to parse input string
             import ast
             try:
-                action_dict = ast.literal_eval(action_input)
+                action_dict = ast.literal_eval(action_input_replaced)
                 if not isinstance(action_dict, dict):
                     action_dict = None
             except Exception as e:
+                logger.error(f"Error parsing action input: {action_input} -> {action_input_replaced}\n{e}")
                 pass
         return AgentAction(
             tool=action, 

bioguider/agents/collection_execute_step.py

@@ -144,7 +144,7 @@ class CollectionExecuteStep(PEOCommonStep):
         agent_executor = AgentExecutor(
             agent=agent,
             tools=self.custom_tools,
-            max_iterations=10,
+            max_iterations=30,
         )
         response = agent_executor.invoke(
             input={"plan_actions": plan_actions, "input": "Now, let's begin."},

bioguider/agents/common_conversation.py

@@ -19,8 +19,26 @@ class CommonConversation:
             callbacks=[callback_handler]
         )
         response = result.generations[0][0].text
-        token_usage = result.llm_output.get("token_usage")
-        return response, vars(token_usage)
+        # Try to normalize token usage across providers
+        token_usage = {}
+        try:
+            if hasattr(result, "llm_output") and result.llm_output is not None:
+                raw = result.llm_output.get("token_usage") or result.llm_output.get("usage")
+                if isinstance(raw, dict):
+                    token_usage = {
+                        "total_tokens": raw.get("total_tokens") or raw.get("total"),
+                        "prompt_tokens": raw.get("prompt_tokens") or raw.get("prompt"),
+                        "completion_tokens": raw.get("completion_tokens") or raw.get("completion"),
+                    }
+        except Exception:
+            pass
+        if not token_usage:
+            token_usage = {
+                "total_tokens": getattr(callback_handler, "total_tokens", 0),
+                "prompt_tokens": getattr(callback_handler, "prompt_tokens", 0),
+                "completion_tokens": getattr(callback_handler, "completion_tokens", 0),
+            }
+        return response, token_usage
     
     def generate_with_schema(self, system_prompt: str, instruction_prompt: str, schema: any):
         system_prompt = escape_braces(system_prompt)

bioguider/agents/consistency_collection_execute_step.py

@@ -0,0 +1,152 @@
+import logging
+
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from langchain.tools import BaseTool, StructuredTool
+from langchain.agents import AgentExecutor, create_react_agent
+from langchain_community.callbacks.openai_info import OpenAICallbackHandler
+
+from bioguider.agents.consistency_collection_task_utils import ConsistencyCollectionWorkflowState
+from bioguider.agents.consistency_collection_task_utils import ConsistencyCollectionWorkflowState
+from bioguider.database.code_structure_db import CodeStructureDb
+from bioguider.utils.constants import DEFAULT_TOKEN_USAGE
+from bioguider.agents.agent_utils import CustomOutputParser, CustomPromptTemplate
+from bioguider.agents.peo_common_step import ( 
+    PEOCommonStep,
+)
+
+logger = logging.getLogger(__name__)
+
+CONSISTENCY_EVAL_EXECUTION_SYSTEM_PROMPT = """You are an expert developer specializing in the biomedical domain.
+
+You are given a **plan** and are expected to complete it using the available tools.
+
+---
+
+### **Available Tools**
+{tools}
+
+---
+
+### **Your Task**
+
+Your job is to **execute the given plan step by step**, using the tools available to you.
+
+---
+
+### **Output Format (Strict Order Required)**
+
+For **each step**, follow the **exact format** below and **do not change the order of the fields** under any circumstances:
+
+```
+Thought: Describe what you are thinking or planning to do next.  
+Action: The tool you are going to use (must be one of: {tool_names})  
+Action Input: The input provided to the selected action  
+Observation: The result returned by the action
+```
+
+---
+
+### **Important Instructions**
+1. You may repeat the **Thought → Action → Action Input → Observation** loop as needed.  
+2. Once all steps in the plan have been executed, output all the results using this format:
+3. For each step, **only execute one tool**.
+
+```
+Thought: I have completed the plan.
+Final Answer:
+Action: <tool_name>
+Action Input: <input>
+Action Observation: <Observation1>
+---
+Action: <tool_name>
+Action Input: <input>
+Action Observation: <Observation2>
+---
+...
+```
+
+---
+
+### **Plan**
+{plan_actions}
+
+### **Actions Already Taken**
+{agent_scratchpad}
+
+---
+
+{input}
+"""
+
+class ConsistencyCollectionExecuteStep(PEOCommonStep):
+    def __init__(
+        self, 
+        llm: BaseChatOpenAI, 
+        code_structure_db: CodeStructureDb,
+        custom_tools: list[BaseTool] | None = None,
+    ):
+        self.llm = llm
+        self.code_structure_db = code_structure_db
+        self.step_name = "Consistency Collection Execute Step"
+        self.custom_tools = custom_tools if custom_tools is not None else []
+
+    def _execute_directly(self, state: ConsistencyCollectionWorkflowState):
+        plan_actions = state["plan_actions"]
+        prompt = CustomPromptTemplate(
+            template=CONSISTENCY_EVAL_EXECUTION_SYSTEM_PROMPT,
+            tools=self.custom_tools,
+            plan_actions=plan_actions,
+            input_variables=[
+                "tools", "tool_names", "agent_scratchpad",
+                "intermediate_steps", "plan_actions",
+            ],
+        )
+        output_parser = CustomOutputParser()
+        agent = create_react_agent(
+            llm=self.llm,
+            tools=self.custom_tools,
+            prompt=prompt,
+            output_parser=output_parser,
+            stop_sequence=["\nObservation:"],
+        )
+        callback_handler = OpenAICallbackHandler()
+        agent_executor = AgentExecutor(
+            agent=agent,
+            tools=self.custom_tools,
+            max_iterations=30,
+        )
+        response = agent_executor.invoke(
+            input={
+                "plan_actions": plan_actions,
+                "input": "Now, let's begin."
+            },
+            config={
+                "callbacks": [callback_handler],
+                "recursion_limit": 20,
+            },
+        )
+        if "output" in response:
+            output = response["output"]
+            self._print_step(state, step_output=f"**Execute Output:** \n{output}")
+            if "**Final Answer**" in output:
+                final_answer = output.split("**Final Answer:**")[-1].strip().strip(":")
+                step_output = final_answer
+            elif "Final Answer" in output:
+                final_answer = output.split("Final Answer")[-1].strip().strip(":")
+                step_output = final_answer
+            else:
+                step_output = output
+            self._print_step(state, step_output=step_output)
+            state["step_output"] = step_output
+        else:
+            logger.error("No output found in the response.")
+            self._print_step(
+                state,
+                step_output="Error: No output found in the response.",
+            )
+            state["step_output"] = "Error: No output found in the response."
+        
+        token_usage = vars(callback_handler)
+        token_usage = {**DEFAULT_TOKEN_USAGE, **token_usage}
+            
+        return state, token_usage

bioguider/agents/consistency_collection_observe_step.py

@@ -0,0 +1,128 @@
+from typing import Callable
+
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from langchain_core.prompts import ChatPromptTemplate
+from bioguider.agents.agent_utils import ObservationResult
+from bioguider.agents.common_agent_2step import CommonAgentTwoSteps
+from bioguider.agents.consistency_collection_task_utils import ConsistencyCollectionWorkflowState
+from bioguider.agents.peo_common_step import PEOCommonStep
+
+CONSISTENCY_EVAL_OBSERVE_SYSTEM_PROMPT = """You are an expert developer specializing in the biomedical domain.
+
+### **Goal**
+Your task is to collect the function, class, and method definitions and docstrings for a given user guide/API documentation.
+
+---
+
+### **Intermediate Steps**  
+Here are the results from previous steps:  
+{intermediate_steps}
+
+---
+
+### **Instructions**
+1. Your goal is if you have enough information to evaluate the consistency of the user guide/API documentation.
+2. Carefully review the **Goal**, **User Guide/API Documentation**, and **Intermediate Output**.
+3. If you believe you have enough information to evaluate the consistency of the user guide/API documentation:
+
+* Proceed with the following format:
+
+  * Provide your reasoning under **Analysis**
+  * Then give your final answer under **FinalAnswer**
+  * **FinalAnswer** format must exactly match this format:
+    **FinalAnswer**: {{"final_answer": "yes" or "no"}}
+  * Your answer **must exactly match the follwing format** (note: no JSON code block, no additional comments), **do not** make up anything:
+
+  ```
+  **Analysis**: your analysis here 
+  **FinalAnswer**: {{"final_answer": "yes" or "no"}}
+  ```
+4. If you believe you do not have enough information to evaluate the consistency of the user guide/API documentation:
+
+* Provide your reasoning under **Thoughts**:
+
+  ```
+  **Thoughts**: your thoughts here
+  ```
+
+Be precise and support your reasoning with evidence from the input.
+
+---
+
+### **Notes**
+We are collecting information over multiple rounds, your thoughts and the output of this step will be persisted, so please **do not rush to provide a Final Answer**.  
+If you find the current information insufficient, share your reasoning or thoughts instead—we’ll continue with the next round accordingly.
+
+---
+
+### **Input User Guide/API Documentation**
+{user_guide_api_documentation}
+
+---
+
+"""
+
+
+class ConsistencyCollectionObserveStep(PEOCommonStep):
+    def __init__(
+        self,
+        llm: BaseChatOpenAI,
+    ):
+        super().__init__(llm=llm)
+        self.step_name = "Consistency Collection Observe Step"
+
+    def _build_prompt(self, state):
+        user_guide_api_documentation = state["user_guide_api_documentation"]
+        intermediate_steps = self._build_intermediate_steps(state)
+        prompt = ChatPromptTemplate.from_template(CONSISTENCY_EVAL_OBSERVE_SYSTEM_PROMPT)
+        return prompt.format(
+            user_guide_api_documentation=user_guide_api_documentation,
+            intermediate_steps=intermediate_steps,
+        )
+
+    def _collect_final_answer(self, state: ConsistencyCollectionWorkflowState):
+        if not ("final_answer" in state and state["final_answer"] is not None and
+            state["final_answer"].strip().lower() == "yes"):
+            return None
+        
+        final_result = ""
+        if "intermediate_steps" in state and state["intermediate_steps"] is not None:
+            for i in range(len(state["intermediate_steps"])):
+                final_result += state["intermediate_steps"][i]
+                final_result += "\n\n"
+        if "step_output" in state and state["step_output"] is not None:
+            final_result += state["step_output"]
+            final_result += "\n\n"
+        
+        return final_result
+
+            
+    def _execute_directly(self, state: ConsistencyCollectionWorkflowState):
+        system_prompt = self._build_prompt(state)
+        agent = CommonAgentTwoSteps(llm=self.llm)
+        res, _, token_usage, reasoning_process = agent.go(
+            system_prompt=system_prompt,
+            instruction_prompt="Now, let's begin the consistency collection observe step.",
+            schema=ObservationResult,
+        )
+        state["final_answer"] = res.FinalAnswer
+        analysis = res.Analysis
+        thoughts = res.Thoughts
+        state["step_analysis"] = analysis
+        state["step_thoughts"] = thoughts
+        state["step_count"] += 1
+        state["final_assembly_result"] = self._collect_final_answer(state)
+        self._print_step(
+            state,
+            step_output=f"**Observation Reasoning Process {state['step_count']}**\n{reasoning_process}"
+        )
+        self._print_step(
+            state,
+            step_output=f"Final Answer: {res.FinalAnswer if res.FinalAnswer else None}\nAnalysis: {analysis}\nThoughts: {thoughts}",
+        )
+        if state["final_assembly_result"] is not None:
+            self._print_step(
+                state,
+                step_output=f"Final Assembly Result: {state['final_assembly_result']}",
+            )
+        return state, token_usage

bioguider/agents/consistency_collection_plan_step.py

@@ -0,0 +1,128 @@
+from langchain_openai.chat_models.base import BaseChatOpenAI
+
+from langchain.tools import BaseTool
+from langchain_core.prompts import ChatPromptTemplate, StringPromptTemplate
+from bioguider.agents.agent_utils import (
+    convert_plan_to_string, 
+    get_tool_names_and_descriptions,
+    PlanAgentResultJsonSchema,
+    PlanAgentResult,
+)
+from bioguider.agents.common_agent_2step import CommonAgentTwoChainSteps, CommonAgentTwoSteps
+from bioguider.agents.consistency_collection_task_utils import ConsistencyCollectionWorkflowState
+from bioguider.agents.peo_common_step import PEOCommonStep
+
+CONSISTANCE_EVAL_PLAN_SYSTEM_PROMPT = ChatPromptTemplate.from_template("""### **Goal**  
+You are an expert developer specializing in the biomedical domain. 
+Your task is to collect the function, class, and method definitions and docstrings for a given user guide/API documentation.
+
+---
+
+### **Function Tools**  
+You have access to the following function tools:  
+{tools}
+
+---
+
+### **Intermediate Steps**  
+Here are the results from previous steps:  
+{intermediate_steps}
+
+---
+
+### **Intermediate Thoughts**  
+- **Analysis**: {intermediate_analysis}  
+- **Thoughts**: {intermediate_thoughts}
+
+---
+
+### **Instructions**
+1. We will iterate through multiple **Plan -> Execution -> Observation** loops as needed.  
+   - All variables and tool outputs are **persisted across rounds**, so you can build on prior results.  
+   - Develop your plan **incrementally**, and reflect on intermediate observations before proceeding.  
+   - Limit each step to **one or two actions** — avoid trying to complete everything in a single step.
+
+2. Your task is to evaluate the consistency of the user guide/API documentation.
+
+3. You may use 
+   - the `retrieve_function_definition_and_docstring` tool to get the function definition and docstring or,
+   - the `retrieve_class_definition_and_docstring` to get the class definition and docstring or,
+   - the `retrieve_class_and_method_definition_and_docstring` to get the class and method definition and docstring.
+
+4. Your plan can only use the above tools, **do not** make up any tools not in the above tools list.
+
+5. If no function, class, or method is found in the given user guide/API documentation, you should return "N/A" as an empty plan.
+   Our tools can only retrieve the **function**, **class**, **method** definition and docstring, **do not** make up any function, class, or method name.
+    
+    
+### **Input User Guide/API Documentation**
+{user_guide_api_documentation}
+
+### **Output Format**
+Your plan **must exactly match** a sequence of steps in the following format, **do not** make up anything:
+
+Step: <tool name>   # Tool name **must be one** of {tool_names}  
+Step Input: <function/class/method name>
+Step Input: <file path, if not sure, just put "N/A">
+
+Step: <tool name>  # Tool name **must be one** of {tool_names}  
+Step Input: <function/class/method name>
+Step Input: <file path, if not sure, just put "N/A">
+...
+
+...
+""")
+
+class ConsistencyCollectionPlanStep(PEOCommonStep):
+    """
+    ConsistencyCollectionPlanStep is a step in the consistency collection plan process.
+    It is responsible for initializing the tools and compiling the step.
+    """
+    
+    def __init__(
+        self, 
+        llm: BaseChatOpenAI,
+        custom_tools: list[BaseTool] | None = None,
+    ):
+        super().__init__(llm)
+        self.step_name = "Consistency Collection Plan Step"
+        self.custom_tools = custom_tools if custom_tools is not None else []
+    
+    def _prepare_system_prompt(self, state: ConsistencyCollectionWorkflowState) -> str:
+        user_guide_api_documentation = state["user_guide_api_documentation"]
+        intermediate_steps = self._build_intermediate_steps(state)
+        step_analysis, step_thoughts = self._build_intermediate_analysis_and_thoughts(state)
+        tool_names, tools_desc = get_tool_names_and_descriptions(self.custom_tools)
+        system_prompt = CONSISTANCE_EVAL_PLAN_SYSTEM_PROMPT.format(
+            tools=tools_desc,
+            intermediate_steps=intermediate_steps,
+            intermediate_analysis=step_analysis,
+            intermediate_thoughts=step_thoughts,
+            tool_names=tool_names,
+            user_guide_api_documentation=user_guide_api_documentation,
+        )
+        self._print_step(
+            state,
+            step_output="**Intermediate Step Output**\n" + intermediate_steps
+        )
+        self._print_step(
+            state,
+            step_output="**Intermediate Step Analysis**\n{step_analysis}\n**Intermediate Step Thoughts**\n{step_thoughts}",
+        )
+        return system_prompt
+
+    def _execute_directly(self, state: ConsistencyCollectionWorkflowState):
+        system_prompt = self._prepare_system_prompt(state)
+        agent = CommonAgentTwoSteps(llm=self.llm)
+        res, _, token_usage, reasoning_process = agent.go(
+            system_prompt=system_prompt,
+            instruction_prompt="Now, let's begin the consistency collection plan step.",
+            schema=PlanAgentResultJsonSchema,
+        )
+        PEOCommonStep._reset_step_state(state)
+        res = PlanAgentResult(**res)
+        self._print_step(state, step_output=f"**Reasoning Process**\n{reasoning_process}")
+        self._print_step(state, step_output=f"**Plan**\n{str(res.actions)}")
+        state["plan_actions"] = convert_plan_to_string(res)
+
+        return state, token_usage

bioguider/agents/consistency_collection_task.py

@@ -0,0 +1,109 @@
+import os
+from typing import Callable, Optional, TypedDict
+from langchain.prompts import ChatPromptTemplate
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from langchain_core.messages import AIMessage
+from langchain.tools import StructuredTool
+from pydantic import BaseModel, Field
+import logging
+from langgraph.graph import StateGraph, START, END
+
+from bioguider.agents.agent_task import AgentTask
+from bioguider.agents.agent_tools import agent_tool
+from bioguider.agents.agent_utils import read_file, summarize_file
+from bioguider.agents.peo_common_step import PEOWorkflowState
+from bioguider.agents.common_agent import CommonAgent
+from bioguider.agents.common_agent_2step import CommonAgentTwoSteps
+from bioguider.agents.consistency_collection_task_utils import (
+    ConsistencyCollectionWorkflowState,
+    retrieve_function_definition_and_docstring_tool,
+    retrieve_class_definition_and_docstring_tool,
+    retrieve_class_and_method_definition_and_docstring_tool,
+    retrieve_method_definition_and_docstring_tool,
+)
+from bioguider.agents.consistency_collection_plan_step import ConsistencyCollectionPlanStep
+from bioguider.agents.consistency_collection_observe_step import ConsistencyCollectionObserveStep
+from bioguider.agents.consistency_collection_execute_step import ConsistencyCollectionExecuteStep
+from bioguider.database.code_structure_db import CodeStructureDb
+
+logger = logging.getLogger(__name__)
+
+class ConsistencyCollectionTask(AgentTask):
+    def __init__(
+        self, 
+        llm: BaseChatOpenAI,
+        code_structure_db: CodeStructureDb,
+        step_callback: Callable | None = None,
+    ):
+        super().__init__(llm=llm, step_callback=step_callback)
+        self.llm = llm
+        self.code_structure_db = code_structure_db
+        
+        func_tool = retrieve_function_definition_and_docstring_tool(llm=llm, code_structure_db=code_structure_db)
+        class_tool = retrieve_class_definition_and_docstring_tool(llm=llm, code_structure_db=code_structure_db)
+        class_and_method_tool = retrieve_class_and_method_definition_and_docstring_tool(llm=llm, code_structure_db=code_structure_db)
+        method_tool = retrieve_method_definition_and_docstring_tool(llm=llm, code_structure_db=code_structure_db)
+        self.tools = [func_tool, class_tool, class_and_method_tool, method_tool]
+        self.custom_tools = [            
+            StructuredTool.from_function(
+                func_tool.run,
+                description=func_tool.__class__.__doc__,
+                name=func_tool.__class__.__name__,
+            ),            
+            StructuredTool.from_function(
+                class_tool.run,
+                description=class_tool.__class__.__doc__,
+                name=class_tool.__class__.__name__,
+            ),            
+            StructuredTool.from_function(
+                class_and_method_tool.run,
+                description=class_and_method_tool.__class__.__doc__,
+                name=class_and_method_tool.__class__.__name__,
+            ),
+            StructuredTool.from_function(
+                method_tool.run,
+                description=method_tool.__class__.__doc__,
+                name=method_tool.__class__.__name__,
+            ),
+        ]
+
+        self.steps = [
+            ConsistencyCollectionPlanStep(llm=llm, custom_tools=self.custom_tools), 
+            ConsistencyCollectionExecuteStep(llm=llm, code_structure_db=code_structure_db, custom_tools=self.custom_tools),
+            ConsistencyCollectionObserveStep(llm=llm)
+        ]
+
+    def _compile(self, repo_path: str, gitignore_path: str):
+        def check_observe_step(state: ConsistencyCollectionWorkflowState):
+            if "final_answer" in state and state["final_answer"] is not None:
+                return END
+            return "plan_step"
+        def check_plan_step(state: ConsistencyCollectionWorkflowState):
+            if "plan_actions" in state and state["plan_actions"] is not None and len(state["plan_actions"]) > 0:
+                return "execute_step"
+            return END
+        
+        graph = StateGraph(ConsistencyCollectionWorkflowState)
+        graph.add_node("plan_step", self.steps[0].execute)
+        graph.add_node("execute_step", self.steps[1].execute)
+        graph.add_node("observe_step", self.steps[2].execute)
+        graph.add_edge(START, "plan_step")
+        graph.add_conditional_edges("plan_step", check_plan_step, {"execute_step", END})
+        graph.add_edge("execute_step", "observe_step")
+        graph.add_conditional_edges("observe_step", check_observe_step, {"plan_step", END})
+
+        self.graph = graph.compile()
+
+    def collect(self, user_guide_api_documentation: str) -> tuple[bool, str | None]:
+        s = self._go_graph({
+            "user_guide_api_documentation": user_guide_api_documentation,
+            "step_count": 0,
+        })
+        # analyze the final assembly result
+        if "final_assembly_result" in s and s["final_assembly_result"] is not None:
+            self._print_step(step_name="Final Assembly Result")
+            self._print_step(step_output=s["final_assembly_result"])
+            return True, s["final_assembly_result"]
+        else:
+            return False, s["thoughts"] if "thoughts" in s else None
+