From e8c55c349fda77e86347b300f4d887b63430f8aa Mon Sep 17 00:00:00 2001 From: Boris Bergsma Date: Mon, 16 Sep 2024 15:59:57 +0200 Subject: [PATCH] fix all weird bugs --- CHANGELOG.md | 5 +++- pyproject.toml | 1 + src/neuroagent/agents/base_agent.py | 11 +++----- src/neuroagent/agents/simple_agent.py | 18 ++++++------- src/neuroagent/agents/simple_chat_agent.py | 22 +++++++--------- .../multi_agents/base_multi_agent.py | 6 ++--- .../multi_agents/supervisor_multi_agent.py | 20 +++++++------- src/neuroagent/tools/base_tool.py | 26 +++++++++---------- src/neuroagent/tools/electrophys_tool.py | 2 +- src/neuroagent/tools/get_morpho_tool.py | 2 +- .../tools/kg_morpho_features_tool.py | 13 +++++----- .../tools/literature_search_tool.py | 2 +- .../tools/morphology_features_tool.py | 2 +- .../tools/resolve_brain_region_tool.py | 2 +- src/neuroagent/tools/traces_tool.py | 2 +- tests/tools/test_basic_tool.py | 11 ++++---- 16 files changed, 69 insertions(+), 76 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9404315..fd91bc3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,5 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Removed +- Github action to create the docs. - +### Changed +- Migration to pydantic V2. diff --git a/pyproject.toml b/pyproject.toml index d7c2102..5fc5406 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "langgraph-checkpoint-postgres", "langgraph-checkpoint-sqlite", "neurom", + "psycopg-binary", "psycopg2-binary", "pydantic-settings", "python-dotenv", diff --git a/src/neuroagent/agents/base_agent.py b/src/neuroagent/agents/base_agent.py index 13230ef..df0833d 100644 --- a/src/neuroagent/agents/base_agent.py +++ b/src/neuroagent/agents/base_agent.py @@ -4,7 +4,6 @@ from typing import Any, AsyncIterator from langchain.chat_models.base import BaseChatModel -from langchain.llms.base import BaseLLM from langchain_core.messages import ( AIMessage, ChatMessage, @@ -21,9 +20,7 @@ SystemMessagePromptTemplate, ) from langchain_core.tools import BaseTool -from pydantic import BaseModel as BaseModelV2 -from pydantic import ConfigDict -from pydantic.v1 import BaseModel +from pydantic import BaseModel, ConfigDict BASE_PROMPT = ChatPromptTemplate( input_variables=["agent_scratchpad", "input"], @@ -63,14 +60,14 @@ ) -class AgentStep(BaseModelV2): +class AgentStep(BaseModel): """Class for agent decision steps.""" tool_name: str arguments: dict[str, Any] | str -class AgentOutput(BaseModelV2): +class AgentOutput(BaseModel): """Class for agent response.""" response: str @@ -81,7 +78,7 @@ class AgentOutput(BaseModelV2): class BaseAgent(BaseModel, ABC): """Base class for services.""" - llm: BaseLLM | BaseChatModel + llm: BaseChatModel tools: list[BaseTool] agent: Any diff --git a/src/neuroagent/agents/simple_agent.py b/src/neuroagent/agents/simple_agent.py index 8c4214f..819bd38 100644 --- a/src/neuroagent/agents/simple_agent.py +++ b/src/neuroagent/agents/simple_agent.py @@ -5,8 +5,7 @@ from langchain_core.messages import AIMessage from langgraph.prebuilt import create_react_agent -from pydantic import ConfigDict -from pydantic.v1 import root_validator +from pydantic import model_validator from neuroagent.agents import AgentOutput, AgentStep, BaseAgent @@ -16,20 +15,19 @@ class SimpleAgent(BaseAgent): """Simple Agent class.""" - model_config = ConfigDict(arbitrary_types_allowed=True) - - @root_validator(pre=True) - def create_agent(cls, values: dict[str, Any]) -> dict[str, Any]: + @model_validator(mode="before") + @classmethod + def create_agent(cls, data: dict[str, Any]) -> dict[str, Any]: """Instantiate the clients upon class creation.""" # Initialise the agent with the tools - values["agent"] = create_react_agent( - model=values["llm"], - tools=values["tools"], + data["agent"] = create_react_agent( + model=data["llm"], + tools=data["tools"], state_modifier="""You are a helpful assistant helping scientists with neuro-scientific questions. You must always specify in your answers from which brain regions the information is extracted. Do no blindly repeat the brain region requested by the user, use the output of the tools instead.""", ) - return values + return data def run(self, query: str) -> Any: """Run the agent against a query. diff --git a/src/neuroagent/agents/simple_chat_agent.py b/src/neuroagent/agents/simple_chat_agent.py index f706b51..1262a2c 100644 --- a/src/neuroagent/agents/simple_chat_agent.py +++ b/src/neuroagent/agents/simple_chat_agent.py @@ -6,7 +6,7 @@ from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.prebuilt import create_react_agent -from pydantic.v1 import root_validator +from pydantic import model_validator from neuroagent.agents import AgentOutput, AgentStep, BaseAgent @@ -18,23 +18,19 @@ class SimpleChatAgent(BaseAgent): memory: BaseCheckpointSaver - class Config: - """Config.""" - - arbitrary_types_allowed = True - - @root_validator(pre=True) - def create_agent(cls, values: dict[str, Any]) -> dict[str, Any]: + @model_validator(mode="before") + @classmethod + def create_agent(cls, data: dict[str, Any]) -> dict[str, Any]: """Instantiate the clients upon class creation.""" - values["agent"] = create_react_agent( - model=values["llm"], - tools=values["tools"], - checkpointer=values["memory"], + data["agent"] = create_react_agent( + model=data["llm"], + tools=data["tools"], + checkpointer=data["memory"], state_modifier="""You are a helpful assistant helping scientists with neuro-scientific questions. You must always specify in your answers from which brain regions the information is extracted. Do no blindly repeat the brain region requested by the user, use the output of the tools instead.""", ) - return values + return data def run(self, session_id: str, query: str) -> Any: """Run the agent against a query.""" diff --git a/src/neuroagent/multi_agents/base_multi_agent.py b/src/neuroagent/multi_agents/base_multi_agent.py index ab169b8..8a3f62a 100644 --- a/src/neuroagent/multi_agents/base_multi_agent.py +++ b/src/neuroagent/multi_agents/base_multi_agent.py @@ -4,9 +4,7 @@ from typing import Any, AsyncIterator from langchain.chat_models.base import BaseChatModel -from langchain.llms.base import BaseLLM -from pydantic import ConfigDict -from pydantic.v1 import BaseModel +from pydantic import BaseModel, ConfigDict from neuroagent.agents import AgentOutput from neuroagent.tools.base_tool import BasicTool @@ -15,7 +13,7 @@ class BaseMultiAgent(BaseModel, ABC): """Base class for multi agents.""" - llm: BaseLLM | BaseChatModel + llm: BaseChatModel main_agent: Any agents: list[tuple[str, list[BasicTool]]] diff --git a/src/neuroagent/multi_agents/supervisor_multi_agent.py b/src/neuroagent/multi_agents/supervisor_multi_agent.py index 5d067bf..15df2e3 100644 --- a/src/neuroagent/multi_agents/supervisor_multi_agent.py +++ b/src/neuroagent/multi_agents/supervisor_multi_agent.py @@ -16,8 +16,7 @@ from langgraph.graph import END, START, StateGraph from langgraph.graph.graph import CompiledGraph from langgraph.prebuilt import create_react_agent -from pydantic import ConfigDict -from pydantic.v1 import root_validator +from pydantic import ConfigDict, model_validator from neuroagent.agents import AgentOutput, AgentStep from neuroagent.multi_agents.base_multi_agent import BaseMultiAgent @@ -39,8 +38,9 @@ class SupervisorMultiAgent(BaseMultiAgent): model_config = ConfigDict(arbitrary_types_allowed=True) - @root_validator(pre=True) - def create_main_agent(cls, values: dict[str, Any]) -> dict[str, Any]: + @model_validator(mode="before") + @classmethod + def create_main_agent(cls, data: dict[str, Any]) -> dict[str, Any]: """Instantiate the clients upon class creation.""" logger.info("Creating main agent, supervisor and all the agents with tools.") system_prompt = ( @@ -50,7 +50,7 @@ def create_main_agent(cls, values: dict[str, Any]) -> dict[str, Any]: " task and respond with their results and status. When finished," " respond with FINISH." ) - agents_list = [elem[0] for elem in values["agents"]] + agents_list = [elem[0] for elem in data["agents"]] logger.info(f"List of agents name: {agents_list}") options = ["FINISH"] + agents_list @@ -84,14 +84,14 @@ def create_main_agent(cls, values: dict[str, Any]) -> dict[str, Any]: ), ] ).partial(options=str(options), members=", ".join(agents_list)) - values["main_agent"] = ( + data["main_agent"] = ( prompt - | values["llm"].bind_functions( + | data["llm"].bind_functions( functions=[function_def], function_call="route" ) | JsonOutputFunctionsParser() ) - values["summarizer"] = ( + data["summarizer"] = ( PromptTemplate.from_template( """You are an helpful assistant. Here is the question of the user: {question}. And here are the results of the different tools used to answer: {responses}. @@ -101,10 +101,10 @@ def create_main_agent(cls, values: dict[str, Any]) -> dict[str, Any]: Please formulate a complete response to give to the user ONLY based on the results. """ ) - | values["llm"] + | data["llm"] ) - return values + return data @staticmethod async def agent_node( diff --git a/src/neuroagent/tools/base_tool.py b/src/neuroagent/tools/base_tool.py index ad6041f..6bb9003 100644 --- a/src/neuroagent/tools/base_tool.py +++ b/src/neuroagent/tools/base_tool.py @@ -4,9 +4,8 @@ import logging from typing import Any -from langchain_core.pydantic_v1 import ValidationError from langchain_core.tools import BaseTool, ToolException -from pydantic.v1 import BaseModel, root_validator +from pydantic import BaseModel, ValidationError, model_validator logger = logging.getLogger(__name__) @@ -14,23 +13,21 @@ def process_validation_error(error: ValidationError) -> str: """Handle validation errors when tool inputs are wrong.""" error_list = [] - - # not happy with this solution but it is to extract the name of the input class - name = str(error.model).split(".")[-1].strip(">") + name = error.title # We have to iterate, in case there are multiple errors. try: for err in error.errors(): - if "ctx" in err: + if err["type"] == "literal_error": error_list.append( { "Validation error": ( - f'Wrong value: {err["ctx"]["given"]} for input' + f'Wrong value: provided {err["input"]} for input' f' {err["loc"][0]}. Try again and change this problematic' " input." ) } ) - elif "loc" in err and err["msg"] == "field required": + elif err["type"] == "missing": error_list.append( { "Validation error": ( @@ -70,12 +67,13 @@ class BasicTool(BaseTool): name: str = "base" description: str = "Base tool from which regular tools should inherit." - @root_validator(pre=True) - def handle_errors(cls, values: dict[str, Any]) -> dict[str, Any]: + @model_validator(mode="before") + @classmethod + def handle_errors(cls, data: dict[str, Any]) -> dict[str, Any]: """Instantiate the clients upon class creation.""" - values["handle_validation_error"] = process_validation_error - values["handle_tool_error"] = process_tool_error - return values + data["handle_validation_error"] = process_validation_error + data["handle_tool_error"] = process_tool_error + return data class BaseToolOutput(BaseModel): @@ -83,4 +81,4 @@ class BaseToolOutput(BaseModel): def __repr__(self) -> str: """Representation method.""" - return self.json() + return self.model_dump_json() diff --git a/src/neuroagent/tools/electrophys_tool.py b/src/neuroagent/tools/electrophys_tool.py index b3f0f3c..e3f7018 100644 --- a/src/neuroagent/tools/electrophys_tool.py +++ b/src/neuroagent/tools/electrophys_tool.py @@ -7,8 +7,8 @@ from bluepyefe.extract import extract_efeatures from efel.units import get_unit -from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.tools import ToolException +from pydantic import BaseModel, Field from neuroagent.tools.base_tool import BaseToolOutput, BasicTool from neuroagent.utils import get_kg_data diff --git a/src/neuroagent/tools/get_morpho_tool.py b/src/neuroagent/tools/get_morpho_tool.py index 0bedd6f..5bd73c4 100644 --- a/src/neuroagent/tools/get_morpho_tool.py +++ b/src/neuroagent/tools/get_morpho_tool.py @@ -3,8 +3,8 @@ import logging from typing import Any, Optional, Type -from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.tools import ToolException +from pydantic import BaseModel, Field from neuroagent.cell_types import get_celltypes_descendants from neuroagent.tools.base_tool import BaseToolOutput, BasicTool diff --git a/src/neuroagent/tools/kg_morpho_features_tool.py b/src/neuroagent/tools/kg_morpho_features_tool.py index 6216427..ba67c0f 100644 --- a/src/neuroagent/tools/kg_morpho_features_tool.py +++ b/src/neuroagent/tools/kg_morpho_features_tool.py @@ -3,8 +3,8 @@ import logging from typing import Any, Literal, Type -from langchain_core.pydantic_v1 import BaseModel, Field, root_validator from langchain_core.tools import ToolException +from pydantic import BaseModel, Field, model_validator from neuroagent.tools.base_tool import BaseToolOutput, BasicTool from neuroagent.utils import get_descendants_id @@ -121,13 +121,14 @@ class FeatureInput(BaseModel): ) feat_range: FeatRangeInput | None = None - @root_validator(pre=True) - def check_if_list(cls, values: Any) -> dict[str, str | list[float | int] | None]: + @model_validator(mode="before") + @classmethod + def check_if_list(cls, data: Any) -> dict[str, str | list[float | int] | None]: """Validate that the values passed to the constructor are a dictionary.""" - if isinstance(values, list) and len(values) == 1: - data_dict = values[0] + if isinstance(data, list) and len(data) == 1: + data_dict = data[0] else: - data_dict = values + data_dict = data return data_dict diff --git a/src/neuroagent/tools/literature_search_tool.py b/src/neuroagent/tools/literature_search_tool.py index 2ae5886..b25d9bb 100644 --- a/src/neuroagent/tools/literature_search_tool.py +++ b/src/neuroagent/tools/literature_search_tool.py @@ -3,8 +3,8 @@ import logging from typing import Any, Type -from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.tools import ToolException +from pydantic import BaseModel, Field from neuroagent.tools.base_tool import BaseToolOutput, BasicTool diff --git a/src/neuroagent/tools/morphology_features_tool.py b/src/neuroagent/tools/morphology_features_tool.py index b2bb2e8..aee5f72 100644 --- a/src/neuroagent/tools/morphology_features_tool.py +++ b/src/neuroagent/tools/morphology_features_tool.py @@ -5,9 +5,9 @@ import neurom import numpy as np -from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.tools import ToolException from neurom.io.utils import load_morphology +from pydantic import BaseModel, Field from neuroagent.tools.base_tool import BaseToolOutput, BasicTool from neuroagent.utils import get_kg_data diff --git a/src/neuroagent/tools/resolve_brain_region_tool.py b/src/neuroagent/tools/resolve_brain_region_tool.py index eec2306..ecd4c4c 100644 --- a/src/neuroagent/tools/resolve_brain_region_tool.py +++ b/src/neuroagent/tools/resolve_brain_region_tool.py @@ -3,8 +3,8 @@ import logging from typing import Any, Optional, Type -from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.tools import ToolException +from pydantic import BaseModel, Field from neuroagent.resolving import resolve_query from neuroagent.tools.base_tool import BaseToolOutput, BasicTool diff --git a/src/neuroagent/tools/traces_tool.py b/src/neuroagent/tools/traces_tool.py index 9b384bd..ce3f30c 100644 --- a/src/neuroagent/tools/traces_tool.py +++ b/src/neuroagent/tools/traces_tool.py @@ -3,8 +3,8 @@ import logging from typing import Any, Literal, Optional, Type -from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.tools import ToolException +from pydantic import BaseModel, Field from neuroagent.tools.base_tool import BaseToolOutput, BasicTool from neuroagent.utils import get_descendants_id diff --git a/tests/tools/test_basic_tool.py b/tests/tools/test_basic_tool.py index 1ee88f4..fb9ab17 100644 --- a/tests/tools/test_basic_tool.py +++ b/tests/tools/test_basic_tool.py @@ -2,16 +2,16 @@ from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel from langchain_core.messages import AIMessage, HumanMessage -from langchain_core.pydantic_v1 import BaseModel from langchain_core.tools import ToolException from langgraph.prebuilt import create_react_agent from neuroagent.tools.base_tool import BasicTool +from pydantic import BaseModel class input_for_test(BaseModel): test_str: str test_int: int - test_litteral: Literal["Allowed_1", "Allowed_2"] = None + test_litteral: Literal["Allowed_1", "Allowed_2"] | None = None class basic_tool_for_test(BasicTool): @@ -87,7 +87,7 @@ def bind_tools(self, functions: list): ) assert ( response["messages"][4].content - == '[{"Validation error": "Wrong value: Forbidden_value for input' + == '[{"Validation error": "Wrong value: provided Forbidden_value for input' ' test_litteral. Try again and change this problematic input."}]' ) assert ( @@ -98,7 +98,8 @@ def bind_tools(self, functions: list): ) assert ( response["messages"][6].content - == '[{"Validation error": "test_str. str type expected"}, {"Validation error":' - ' "test_int. value is not a valid integer"}]' + == '[{"Validation error": "test_str. Input should be a valid string"}, ' + '{"Validation error": "test_int. Input should be a valid integer, ' + 'unable to parse string as an integer"}]' ) assert response["messages"][7].content == "fake answer"