bioguider 0.2.52__py3-none-any.whl

bioguider/agents/common_agent.py

@@ -0,0 +1,137 @@
+
+from typing import Any, Callable, Optional
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from langchain_community.callbacks.openai_info import OpenAICallbackHandler
+from langchain_core.messages import SystemMessage, HumanMessage
+from pydantic import BaseModel, Field
+from tenacity import retry, stop_after_attempt, wait_incrementing
+import logging
+
+from bioguider.utils.utils import escape_braces, increase_token_usage
+
+logger = logging.getLogger(__name__)
+
+class RetryException(Exception):
+    """Exception need to retry"""
+
+    pass
+
+class CommonAgentResult(BaseModel):
+    reasoning_process: str = Field(
+        description="A detailed explanation of the thought process or reasoning steps taken to reach a conclusion."
+    )
+
+class CommonAgent:
+    def __init__(self, llm: BaseChatOpenAI):
+        self.llm = llm
+        self.exception: RetryException | None = None
+        self.token_usage: dict | None = None
+
+    def go(
+        self,
+        system_prompt: str,
+        instruction_prompt: str,
+        schema: any,
+        pre_process: Optional[Callable] = None,
+        post_process: Optional[Callable] = None,
+        **kwargs: Optional[Any],
+    ):
+        """
+        execute agent
+
+        Args:
+        system_prompt str: system prompt
+        instruction_prompt str: user prompt to guide how llm execute agent
+        schema pydantic.BaseModel or json schema: llm output result schema
+        pre_process Callable or None: pre-processor that would be executed before llm.invoke
+        post_process Callable or None: post-processor that would be executed after llm.invoke
+        kwargs None or dict: args for pre_proces and post_process
+
+        Return:
+        (output that comply with input args `schema`)
+        """
+        self._initialize()
+        if pre_process is not None:
+            is_OK = pre_process(**kwargs)
+            if not is_OK:  # skip
+                return
+        
+        return self._invoke_agent(
+            system_prompt,
+            instruction_prompt,
+            schema,
+            post_process,
+            **kwargs,
+        )
+
+    def _initialize(self):
+        self.exception = None
+        self.token_usage = None
+
+    def _process_retryexception_message(
+        self, prompt: ChatPromptTemplate
+    ) -> ChatPromptTemplate:
+        if self.exception is None:
+            return prompt
+
+        existing_messages = prompt.messages
+        updated_messages = existing_messages + [("human", str(self.exception))]
+        self.exception = None
+        updated_prompt = ChatPromptTemplate.from_messages(updated_messages)
+        return updated_prompt
+
+    def _incre_token_usage(self, token_usage):
+        incremental_token_usage = token_usage
+        if not isinstance(token_usage, dict):
+            incremental_token_usage = vars(incremental_token_usage)
+        self.token_usage = increase_token_usage(
+            self.token_usage, incremental_token_usage
+        )
+
+    @retry(
+        stop=stop_after_attempt(5),
+        wait=wait_incrementing(start=1.0, increment=3, max=10),
+    )
+    def _invoke_agent(
+        self,
+        system_prompt: str,
+        instruction_prompt: str,
+        schema: any,
+        post_process: Optional[Callable] = None,
+        **kwargs: Optional[Any],
+    ) -> tuple[Any, Any, dict | None, Any | None]:
+        system_prompt = escape_braces(system_prompt)
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", system_prompt),
+            ("human", instruction_prompt),
+        ])
+        # Initialize the callback handler
+        callback_handler = OpenAICallbackHandler()
+
+        updated_prompt = self._process_retryexception_message(prompt)
+        agent = updated_prompt | self.llm.with_structured_output(schema)
+        try:
+            res = agent.invoke(
+                input={},
+                config={
+                    "callbacks": [callback_handler],
+                },
+            )
+            self._incre_token_usage(callback_handler)
+        except Exception as e:
+            logger.error(str(e))
+            raise e
+        processed_res = res
+        if post_process is not None:
+            try:
+                processed_res = post_process(res, **kwargs)
+            except RetryException as e:
+                logger.error(str(e))
+                self.exception = e
+                raise e
+            except Exception as e:
+                logger.error(str(e))
+                raise e
+        return res, processed_res, self.token_usage, None
+    

bioguider/agents/common_agent_2step.py

@@ -0,0 +1,215 @@
+from typing import Any, Callable, Optional
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from langchain_community.callbacks.openai_info import OpenAICallbackHandler
+from pydantic import BaseModel, Field
+from tenacity import retry, stop_after_attempt, wait_incrementing
+import logging
+
+from bioguider.utils.utils import escape_braces
+from bioguider.agents.common_agent import (
+    CommonAgent,
+    RetryException,
+)
+from bioguider.agents.prompt_utils import COT_USER_INSTRUCTION
+
+logger = logging.getLogger()
+
+
+class CommonAgentTwoSteps(CommonAgent):
+    def __init__(self, llm: BaseChatOpenAI):
+        super().__init__(llm)
+
+    def _initialize(self):
+        self.exceptions = None
+        self.token_usage = None
+
+    def _get_retryexception_message(self) -> list[tuple[str, str]]:
+        if self.exceptions is None:
+            return None
+        return [("human", str(excp)) for excp in self.exceptions]
+
+    def _build_prompt_for_cot_step(
+        self,
+        system_prompt: str,
+        instruction_prompt: str,
+    ):
+        # system_prompt = system_prompt.replace("{", "{{").replace("}", "}}")
+        system_prompt = escape_braces(system_prompt)
+        instruction_prompt = instruction_prompt.replace("{", "{{").replace("}", "}}")
+        msgs = [("system", system_prompt)]
+        msgs = msgs + [("human", instruction_prompt)]
+        exception_msgs = self._get_retryexception_message()
+        if exception_msgs is not None:
+            msgs = msgs + exception_msgs
+        msgs = msgs + [("human", COT_USER_INSTRUCTION)]
+        return ChatPromptTemplate.from_messages(msgs)
+    
+    def _build_prompt_for_final_step(
+        self,
+        system_prompt: str,
+        cot_msg: str,
+    ):
+        system_prompt = system_prompt.replace("{", "{{").replace("}", "}}")
+        msgs = [("system", system_prompt)]
+        cot_msg = cot_msg.replace("{", "{{").replace("}", "}}")
+        msgs = msgs + [(
+            "human",
+            f"Please review the following step-by-step reasoning and provide the answer based on it: ```{cot_msg}```"
+        )]
+        return ChatPromptTemplate.from_messages(msgs)
+
+    @retry(
+        stop=stop_after_attempt(5),
+        wait=wait_incrementing(start=1.0, increment=3, max=10),
+    )
+    def _invoke_agent(
+        self,
+        system_prompt: str,
+        instruction_prompt: str,
+        schema: any,
+        post_process: Optional[Callable] = None,
+        **kwargs: Optional[Any],
+    ):
+        # Initialize the callback handler
+        callback_handler = OpenAICallbackHandler()
+        cot_prompt = self._build_prompt_for_cot_step(
+            system_prompt=system_prompt, 
+            instruction_prompt=instruction_prompt
+        )
+
+        try:
+            # First, use llm to do CoT
+            msgs = cot_prompt.invoke(input={}).to_messages()
+            
+            cot_res = self.llm.generate(messages=[msgs])
+            reasoning_process = cot_res.generations[0][0].text
+            token_usage = cot_res.llm_output.get("token_usage")
+            cot_tokens = {
+                "total_tokens": token_usage.get("total_tokens", 0),
+                "prompt_tokens": token_usage.get("prompt_tokens", 0),
+                "completion_tokens": token_usage.get("completion_tokens", 0),
+            }
+            self._incre_token_usage(cot_tokens)
+        except Exception as e:
+            logger.error(str(e))
+            raise e
+        
+        # Then use the reasoning process to do the structured output
+        updated_prompt = self._build_prompt_for_final_step(
+            system_prompt=system_prompt,
+            cot_msg=reasoning_process,
+        )
+        agent = updated_prompt | self.llm.with_structured_output(schema)
+        try:
+            res = agent.invoke(
+                input={},
+                config={
+                    "callbacks": [callback_handler],
+                },
+            )
+            self._incre_token_usage(callback_handler)
+        except Exception as e:
+            logger.error(str(e))
+            raise e
+        processed_res = None
+        if post_process is not None:
+            try:
+                processed_res = post_process(res, **kwargs)
+            except RetryException as e:
+                logger.error(str(e))
+                self.exceptions = [e] if self.exceptions is None else self.exceptions + [e]
+                raise e
+            except Exception as e:
+                logger.error(str(e))
+                raise e
+        return res, processed_res, self.token_usage, reasoning_process
+    
+FINAL_STEP_SYSTEM_PROMPTS = ChatPromptTemplate.from_template("""
+---
+
+You will be given a response generated by a LLM, which includes a **step-by-step reasoning process** followed by a clearly marked **final answer**.
+
+### **Your Task:**
+
+Extract and return only the content of the **final answer**.
+
+---
+
+### **Important Instructions:**
+1. Your task is to **extract only the final answer** from the provided reasoning process.
+   **Do not** make any judgments, interpretations, or modifications to the content.
+
+### **Input:**
+
+{llm_response}
+
+---
+""")
+
+class CommonAgentTwoChainSteps(CommonAgentTwoSteps):
+    def __init__(self, llm):
+        super().__init__(llm)
+
+    def _invoke_agent(self, system_prompt, instruction_prompt, schema, post_process = None, **kwargs):
+        # Initialize the callback handler
+        callback_handler = OpenAICallbackHandler()
+        processed_system_prompt = system_prompt.replace("{", "{{").replace("}", "}}")
+        cot_prompt = self._build_prompt_for_cot_step(
+            system_prompt=processed_system_prompt, 
+            instruction_prompt=instruction_prompt
+        )
+
+        try:
+            # First, use llm to do CoT
+            msgs = cot_prompt.invoke(input={}).to_messages()
+            
+            cot_res = self.llm.generate(messages=[msgs])
+            if cot_res is None or cot_res.llm_output is None:
+                raise Exception("llm generate invalid output")
+            reasoning_process = cot_res.generations[0][0].text
+            token_usage: Any = cot_res.llm_output.get("token_usage")
+            cot_tokens = {
+                "total_tokens": token_usage.get("total_tokens", 0),
+                "prompt_tokens": token_usage.get("prompt_tokens", 0),
+                "completion_tokens": token_usage.get("completion_tokens", 0),
+            }
+            self._incre_token_usage(cot_tokens)
+        except Exception as e:
+            logger.error(str(e))
+            raise e
+                
+        try:
+            # Then use the reasoning process to do the structured output
+            processed_reasoning_process = reasoning_process.replace("{", "{{").replace("}", "}}")
+            final_msg = FINAL_STEP_SYSTEM_PROMPTS.format(
+                llm_response=processed_reasoning_process,
+            )
+            msgs = [(
+                "human",
+                final_msg,
+            )]
+            final_prompt = ChatPromptTemplate.from_messages(msgs)
+            agent = final_prompt | self.llm.with_structured_output(schema)
+            res = agent.invoke(
+                input={},
+                config={
+                    "callbacks": [callback_handler],
+                },
+            )
+            self._incre_token_usage(callback_handler)
+        except Exception as e:
+            logger.error(str(e))
+            raise e
+        processed_res = None
+        if post_process is not None:
+            try:
+                processed_res = post_process(res, **kwargs)
+            except RetryException as e:
+                logger.error(str(e))
+                self.exceptions = [e] if self.exceptions is None else self.exceptions + [e]
+                raise e
+            except Exception as e:
+                logger.error(str(e))
+                raise e
+        return res, processed_res, self.token_usage, reasoning_process

bioguider/agents/common_conversation.py

@@ -0,0 +1,61 @@
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from langchain_community.callbacks.openai_info import OpenAICallbackHandler
+from bioguider.utils.utils import escape_braces
+
+class CommonConversation:
+    def __init__(self, llm: BaseChatOpenAI):
+        self.llm = llm
+
+    def generate(self, system_prompt: str, instruction_prompt: str):
+        msgs = [
+            SystemMessage(system_prompt),
+            HumanMessage(instruction_prompt),
+        ]
+        callback_handler = OpenAICallbackHandler()
+        result = self.llm.generate(
+            messages=[msgs],
+            callbacks=[callback_handler]
+        )
+        response = result.generations[0][0].text
+        # 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)
+        instruction_prompt = escape_braces(instruction_prompt)
+        msgs = [
+            SystemMessage(system_prompt),
+            HumanMessage(instruction_prompt),
+        ]
+        msgs_template = ChatPromptTemplate.from_messages(messages=msgs)
+        callback_handler = OpenAICallbackHandler()
+        agent = msgs_template | self.llm.with_structured_output(schema)
+        result = agent.invoke(
+            input={},
+            config={
+                "callbacks": [callback_handler],
+            },
+        )
+        token_usage = vars(callback_handler)
+        return result, token_usage
+

bioguider/agents/common_step.py

@@ -0,0 +1,85 @@
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Optional, TypedDict
+import logging
+from langchain_openai.chat_models.base import BaseChatOpenAI
+
+from bioguider.utils.constants import DEFAULT_TOKEN_USAGE
+
+logger = logging.getLogger(__name__)
+
+class CommonState(TypedDict):
+    """
+    CommonState is a TypedDict that defines the structure of the state
+    used in the CommonStep class.
+    """
+    llm: Optional[BaseChatOpenAI]
+    step_output_callback: Optional[Callable]
+
+class CommonStep(ABC):
+    """
+    CommonStep is a base class for defining common steps in a workflow.
+    It provides methods to execute the step and handle exceptions.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.step_name = ""
+
+    def enter_step(self, state):
+        if state["step_output_callback"] is None:
+            return
+        state["step_output_callback"](
+            step_name=self.step_name, 
+        )
+
+    def leave_step(self, state, token_usage: Optional[dict[str, int]] = None):
+        if state["step_output_callback"] is None:
+            return
+        if token_usage is not None:
+            state["step_output_callback"](token_usage=token_usage)
+
+    def execute(self, state):
+        """
+        Execute the step. This method should be overridden by subclasses.
+        """
+        self.enter_step(state)
+        state, token_usage = self._execute_directly(state)
+        self.leave_step(state, token_usage)
+        return state
+
+    def _print_step(
+        self,
+        state,
+        step_name: str | None = None,
+        step_output: str | None = None,
+        token_usage: dict | object | None = None,
+    ):
+        step_callback = state["step_output_callback"]
+        if step_callback is None:
+            return
+        # convert token_usage to dict
+        if token_usage is not None and not isinstance(token_usage, dict):
+            token_usage = vars(token_usage)
+            # In case token_usage.total_tokens is 0
+            token_usage = { **DEFAULT_TOKEN_USAGE, **token_usage }
+        step_callback(
+            step_name=step_name,
+            step_output=step_output,
+            token_usage=token_usage,
+        )                
+
+    @abstractmethod
+    def _execute_directly(self, state) -> tuple[dict, dict[str, int]]:
+        """
+        Execute the step directly. This method should be overridden by subclasses.
+        Args:
+            state (CommonState): The state of the workflow.
+        Returns:
+            tuple[dict, dict[str, int]]: The updated state and token usage.
+        """
+        pass
+
+
+
+

bioguider/agents/consistency_collection_step.py

@@ -0,0 +1,102 @@
+
+
+
+from langchain.prompts import ChatPromptTemplate
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from pydantic import BaseModel, Field
+from bioguider.agents.common_agent_2step import CommonAgentTwoSteps
+from bioguider.agents.consistency_evaluation_task_utils import ConsistencyEvaluationState
+from bioguider.agents.peo_common_step import PEOCommonStep
+
+
+CONSISTANCY_COLLECTION_SYSTEM_PROMPT = """
+### **Goal**  
+You are an expert developer specializing in the biomedical domain.
+You will be given a {domain} documentation. Your task is to collect all the functions, classes, and methods that the {domain} documentation mentions.
+
+---
+
+### **Input {domain} Documentation**
+{documentation}
+
+### **Output Format**
+The collected functions, classes, and methods **must exactly match** the following format, **do not** make up anything:
+
+```
+name: <function/class/method name>
+file_path: <file path, if not sure, just put "N/A">
+parameters: <parameters, if not sure, just put "N/A">
+parent: <parent name, if it is a class method, put the class name as the parent name, if not sure, just put "N/A">
+
+...
+
+```
+
+---
+
+### **Output Example**
+```
+name: __init__
+file_path: src/agents/common_agent.py
+parameters: llm, step_output_callback, summarized_files_db
+parent: CommonAgent
+
+name: _invoke_agent
+file_path: src/agents/common_agent.py
+parameters: system_prompt, instruction_prompt, schema, post_process
+parent: CommonAgent
+
+...
+```
+
+"""
+
+class ConsistencyCollectionResult(BaseModel):
+    functions_and_classes: list[dict] = Field(description="A list of functions and classes that the documentation mentions")
+
+ConsistencyCollectionResultJsonSchema = {
+  "properties": {
+    "functions_and_classes": {
+      "description": "A list of functions and classes that the documentation mentions",
+      "items": {
+        "type": "object"
+      },
+      "title": "Functions And Classes",
+      "type": "array"
+    }
+  },
+  "required": [
+    "functions_and_classes"
+  ],
+  "title": "ConsistencyCollectionResult",
+  "type": "object"
+}
+
+class ConsistencyCollectionStep(PEOCommonStep):
+    def __init__(self, llm: BaseChatOpenAI):
+        super().__init__(llm)
+        self.step_name = "Consistency Collection Step"
+
+    def _prepare_system_prompt(self, state: ConsistencyEvaluationState) -> str:
+        documentation = state["documentation"]
+        domain = state["domain"]
+        return ChatPromptTemplate.from_template(CONSISTANCY_COLLECTION_SYSTEM_PROMPT).format(
+            domain=domain,
+            documentation=documentation,
+        )
+
+    def _execute_directly(self, state: ConsistencyEvaluationState) -> tuple[dict, dict[str, int]]:
+        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 step.",
+            schema=ConsistencyCollectionResultJsonSchema,
+        )
+        res: ConsistencyCollectionResult = ConsistencyCollectionResult.model_validate(res)
+        state["functions_and_classes"] = res.functions_and_classes
+        self._print_step(state, step_output=f"Consistency Collection Result: {res.functions_and_classes}")
+        self._print_step(state, step_output=f"Consistency Collection Reasoning Process: {reasoning_process}")
+        
+        return state, token_usage
+

bioguider/agents/consistency_evaluation_task.py

@@ -0,0 +1,57 @@
+
+
+
+from typing import Callable
+from langchain_openai.chat_models.base import BaseChatOpenAI
+from pydantic import BaseModel
+
+from bioguider.agents.consistency_evaluation_task_utils import ConsistencyEvaluationState
+from bioguider.database.code_structure_db import CodeStructureDb
+from .consistency_collection_step import ConsistencyCollectionStep
+from .consistency_query_step import ConsistencyQueryStep
+from .consistency_observe_step import ConsistencyObserveStep
+
+class ConsistencyEvaluationResult(BaseModel):
+    score: int
+    assessment: str
+    development: list[str]
+    strengths: list[str]
+
+class ConsistencyEvaluationTask:
+    def __init__(
+        self, 
+        llm: BaseChatOpenAI, 
+        code_structure_db: CodeStructureDb, 
+        step_callback: Callable | None = None
+    ):
+        self.llm = llm
+        self.code_structure_db = code_structure_db
+        self.step_callback = step_callback
+
+    def evaluate(self, domain: str, documentation: str) -> ConsistencyEvaluationResult:
+        collection_step = ConsistencyCollectionStep(llm=self.llm)
+        query_step = ConsistencyQueryStep(code_structure_db=self.code_structure_db)
+        observe_step = ConsistencyObserveStep(llm=self.llm)
+
+        state = ConsistencyEvaluationState(
+            domain=domain,
+            documentation=documentation,
+            step_output_callback=self.step_callback,
+        )
+
+        state = collection_step.execute(state)
+        state = query_step.execute(state)
+        state = observe_step.execute(state)
+
+        score = state["consistency_score"]
+        assessment = state["consistency_assessment"]
+        development = state["consistency_development"]
+        strengths = state["consistency_strengths"]
+
+        return ConsistencyEvaluationResult(
+            score=score,
+            assessment=assessment,
+            development=development,
+            strengths=strengths,
+        )
+

bioguider/agents/consistency_evaluation_task_utils.py

@@ -0,0 +1,14 @@
+
+from typing import Callable, Optional, TypedDict
+
+
+class ConsistencyEvaluationState(TypedDict):
+    domain: str
+    documentation: str
+    step_output_callback: Optional[Callable]
+    functions_and_classes: Optional[list[dict]]
+    all_query_rows: Optional[list[any]]
+    consistency_score: Optional[int]
+    consistency_assessment: Optional[str]
+    consistency_development: Optional[list[str]]
+    consistency_strengths: Optional[list[str]]