Merge pull request #169 from HolobiomicsLab/dev_madina

lfnothias · web-flow · commit 085c97346b35 · 2026-04-24T14:27:07.000+02:00
Dev madina
diff --git a/app/core/agents/enpkg/tool_chemicals.py b/app/core/agents/enpkg/tool_chemicals.py
@@ -16,7 +16,7 @@
 from langchain_openai import OpenAIEmbeddings
 
 
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 from typing import Optional
 
@@ -47,7 +47,7 @@ class ChemicalResolver(BaseTool):
         Dict[str, str]: a dictionary that contains the output chemical name and corresponding NPC Class URI.
     """
 
-    args_schema = ChemicalInput
+    args_schema: type[BaseModel] = ChemicalInput
     csv_data: List[Document] = None
     retriever: Any = None
     openai_key: str = None
diff --git a/app/core/agents/enpkg/tool_smiles.py b/app/core/agents/enpkg/tool_smiles.py
@@ -1,7 +1,7 @@
 import requests
 
 from langchain.tools import BaseTool
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 
 from typing import Optional
 
@@ -36,7 +36,7 @@ class SMILESResolver(BaseTool):
     smiles_string = "CCC12CCCN3C1C4(CC3)C(CC2)NC5=CC=CC=C45"
     inchikey = _run(smiles_string)
     """
-    args_schema = SMILESInput
+    args_schema: type[BaseModel] = SMILESInput
     openai_key: str = None
 
     def __init__(self, openai_key: str = None):
diff --git a/app/core/agents/enpkg/tool_target.py b/app/core/agents/enpkg/tool_target.py
@@ -5,7 +5,7 @@
 
 
 
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 
 from typing import Optional
 
@@ -38,7 +38,7 @@ class TargetResolver(BaseTool):
         str: A string containing the ChEMBLTarget notation.
     """
 
-    args_schema = TargetInput
+    args_schema: type[BaseModel] = TargetInput
     openai_key: str = None
 
     def __init__(self, openai_key: str = None):
diff --git a/app/core/agents/enpkg/tool_taxon.py b/app/core/agents/enpkg/tool_taxon.py
@@ -1,9 +1,9 @@
-from typing import Optional
+from typing import ClassVar, Optional
 
 from SPARQLWrapper import JSON, SPARQLWrapper
 
 
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 
 from typing import Optional
 
@@ -34,10 +34,10 @@ class TaxonResolver(BaseTool):
     Returns:
         str: A string that contains the Wikidata IRI if found, otherwise `None`.
     """
-    args_schema = TaxonInput
+    args_schema: type[BaseModel] = TaxonInput
 
-    ENDPOINT_URL = "https://query.wikidata.org/sparql"
-    PREFIXES = """
+    ENDPOINT_URL: ClassVar[str] = "https://query.wikidata.org/sparql"
+    PREFIXES: ClassVar[str] = """
         PREFIX prov: <http://www.w3.org/ns/prov#>
         PREFIX pr: <http://www.wikidata.org/prop/reference/>
         PREFIX wdt: <http://www.wikidata.org/prop/direct/>
diff --git a/app/core/agents/entry/tool_filesparser.py b/app/core/agents/entry/tool_filesparser.py
@@ -2,18 +2,22 @@
 import pandas as pd
 from datetime import datetime
 from langchain.tools import BaseTool
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from pathlib import Path
 from app.core.session import setup_logger, create_user_session
 
 logger = setup_logger(__name__)
 
+class FileAnalyzerInput(BaseModel):
+    """Input schema for FileAnalyzer - no input fields required."""
+    pass
+
 class FileAnalyzer(BaseTool):
     name: str = "FILE_ANALYZER"
     description: str = """
     Analyzes files in a specified directory and provides a summary of their content.
     """
-    args_schema = BaseModel  # Using BaseModel directly since no specific input fields are necessary
+    args_schema: type[BaseModel] = FileAnalyzerInput
     folder_path: Path = None
     openai_key: str = None
     session_id: str = None
diff --git a/app/core/agents/interpreter/tool_interpreter.py b/app/core/agents/interpreter/tool_interpreter.py
@@ -2,7 +2,7 @@
 
 
 from codeinterpreterapi import CodeInterpreterSession, File, settings
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 
 from typing import Optional
@@ -37,7 +37,7 @@ class Interpreter(BaseTool):
     Returns:
         None: Outputs the response after interpreting files.
     """
-    args_schema = InterpreterInput
+    args_schema: type[BaseModel] = InterpreterInput
     openai_key: str = None
     session_id: str = None
 
diff --git a/app/core/agents/interpreter/tool_spectrum.py b/app/core/agents/interpreter/tool_spectrum.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 from langchain.callbacks.manager import CallbackManagerForToolRun
 from app.core.memory.database_manager import tools_database
@@ -29,7 +29,7 @@ class SpectrumPlotter(BaseTool):
     Returns:
        An url with spectrum plot.
     """
-    args_schema = SpectrumPlotInput
+    args_schema: type[BaseModel] = SpectrumPlotInput
     session_id: str = None
     openai_key: str = None
     def __init__(self, openai_key: str, session_id: str):
diff --git a/app/core/agents/sparql/tool_merge_result.py b/app/core/agents/sparql/tool_merge_result.py
@@ -1,4 +1,4 @@
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 
 from typing import Optional
@@ -40,7 +40,7 @@ class OutputMerger(BaseTool):
         output_file_path (str): The path to the temporary CSV file where the output with common Wikidata IDs will be saved.
 
     """
-    args_schema = MergerInput
+    args_schema: type[BaseModel] = MergerInput
     session_id: str = None
 
     def __init__(self, session_id: str):
diff --git a/app/core/agents/sparql/tool_sparql.py b/app/core/agents/sparql/tool_sparql.py
@@ -10,7 +10,7 @@
 
 from langchain.chains.llm import LLMChain
 from langchain_core.prompts.prompt import PromptTemplate
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 from langchain_community.vectorstores import FAISS
 from langchain_openai import OpenAIEmbeddings
@@ -207,8 +207,8 @@ class SparqlInput(BaseModel):
 
 ##Question-answering against an RDF or OWL graph by generating SPARQL statements.
 class GraphSparqlQAChain(BaseTool):
-    name = "SPARQL_QUERY_RUNNER"
-    description = """
+    name: str = "SPARQL_QUERY_RUNNER"
+    description: str = """
     The agent resolve the user's question by querying the knowledge graph database. 
     The two inputs should be a string containing the user's question and a string containing the resolved entities in the question.
 
@@ -226,10 +226,10 @@ class GraphSparqlQAChain(BaseTool):
 
         """
     verbose: bool = True
-    args_schema = SparqlInput
+    args_schema: type[BaseModel] = SparqlInput
     sparql_generation_select_chain: LLMChain = None
     sparql_improvement_chain: LLMChain = None
-    requires_params = True
+    requires_params: bool = True
     graph: RdfGraph = None
     session_id: str = None
 
diff --git a/app/core/agents/sparql/tool_wikidata_query.py b/app/core/agents/sparql/tool_wikidata_query.py
@@ -3,13 +3,13 @@
 import csv
 import tempfile
 from pathlib import Path
-from typing import List, Optional
+from typing import ClassVar, List, Optional
 
 from langchain_core.tools import tool
 from SPARQLWrapper import JSON, SPARQLWrapper
 
 from app.core.session import setup_logger, create_user_session
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 from langchain.callbacks.manager import (
     CallbackManagerForToolRun,
@@ -33,10 +33,10 @@ class WikidataStructureSearch(BaseTool):
         Returns:
             str: A string that contains the path to the file with Wikidata IRIs if found, otherwise `None`.
         """
-    args_schema = WikidataInput
+    args_schema: type[BaseModel] = WikidataInput
 
-    ENDPOINT_URL = "https://query.wikidata.org/sparql"
-    PREFIXES = """
+    ENDPOINT_URL: ClassVar[str] = "https://query.wikidata.org/sparql"
+    PREFIXES: ClassVar[str] = """
        PREFIX wd: <http://www.wikidata.org/entity/>
        PREFIX wdt: <http://www.wikidata.org/prop/direct/>
     """
diff --git a/app/core/agents/validator/tool_validator.py b/app/core/agents/validator/tool_validator.py
@@ -5,7 +5,7 @@
 from langchain.callbacks.manager import (
     CallbackManagerForToolRun,
 )
-from langchain.pydantic_v1 import BaseModel, Field
+from pydantic import BaseModel, Field
 from langchain.tools import BaseTool
 from typing import Optional
 
@@ -32,7 +32,7 @@ class PlantDatabaseChecker(BaseTool):
     """
 
 
-    args_schema = PlantInput
+    args_schema: type[BaseModel] = PlantInput
 
     def __init__(self):
         super().__init__()
diff --git a/app/core/memory/custom_sqlite_file.py b/app/core/memory/custom_sqlite_file.py
@@ -6,9 +6,9 @@
 import sqlite3
 import threading
 from contextlib import contextmanager
-from typing import Optional
+from typing import ClassVar, Optional
+
 
-from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.utils import ConfigurableFieldSpec
 from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint
@@ -20,7 +20,7 @@
 
 class SqliteCheckpointerSaver(BaseCheckpointSaver):
     # Specify the path to the database file
-    database_path = "langgraph_checkpoint.db"
+    database_path: ClassVar[str] = "langgraph_checkpoint.db"
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)  # Ensure proper Pydantic initialization
diff --git a/app/core/memory/database_manager.py b/app/core/memory/database_manager.py
@@ -1,39 +1,40 @@
-import os
-from app.core.memory.tools_database import SqliteToolsDatabaseManager
-from app.core.memory.custom_sqlite_file import SqliteCheckpointerSaver
-import importlib
-
-def tools_database():
-    """
-    Get the database manager based on the environment variables
-
-    :return: The database manager
-    """
-    if os.getenv('DATABASE_URL') and os.getenv("TOOLS_DATABASE_MANAGER_CLASS"):
-        class_path = os.getenv('TOOLS_DATABASE_MANAGER_CLASS')
-        if class_path:
-            # Dynamically import the module and class based on the environment variable
-            module_name, class_name = class_path.rsplit('.', 1)
-            module = importlib.import_module(module_name)
-            manager_class = getattr(module, class_name)
-            return manager_class()
-    else:
-        return SqliteToolsDatabaseManager()
-
-def memory_database():
-    """
-    Get the database manager based on the environment variables
-
-    :return: The database manager
-    """
-
-    if os.getenv('DATABASE_URL') and os.getenv("MEMORY_DATABASE_MANAGER_CLASS"):
-        class_path = os.getenv('MEMORY_DATABASE_MANAGER_CLASS')
-        if class_path:
-            # Dynamically import the module and class based on the environment variable
-            module_name, class_name = class_path.rsplit('.', 1)
-            module = importlib.import_module(module_name)
-            manager_class = getattr(module, class_name)
-            return manager_class()
-    else:
-        return SqliteCheckpointerSaver()
+import os
+from app.core.memory.tools_database import SqliteToolsDatabaseManager
+import importlib
+
+def tools_database():
+    """
+    Get the database manager based on the environment variables
+
+    :return: The database manager
+    """
+    if os.getenv('DATABASE_URL') and os.getenv("TOOLS_DATABASE_MANAGER_CLASS"):
+        class_path = os.getenv('TOOLS_DATABASE_MANAGER_CLASS')
+        if class_path:
+            # Dynamically import the module and class based on the environment variable
+            module_name, class_name = class_path.rsplit('.', 1)
+            module = importlib.import_module(module_name)
+            manager_class = getattr(module, class_name)
+            return manager_class()
+    else:
+        return SqliteToolsDatabaseManager()
+
+def memory_database():
+    """
+    Get the database manager based on the environment variables
+
+    :return: The database manager
+    """
+
+    if os.getenv('DATABASE_URL') and os.getenv("MEMORY_DATABASE_MANAGER_CLASS"):
+        class_path = os.getenv('MEMORY_DATABASE_MANAGER_CLASS')
+        if class_path:
+            # Dynamically import the module and class based on the environment variable
+            module_name, class_name = class_path.rsplit('.', 1)
+            module = importlib.import_module(module_name)
+            manager_class = getattr(module, class_name)
+            return manager_class()
+    else:
+        # Use langgraph's built-in MemorySaver which implements the current checkpoint API
+        from langgraph.checkpoint.memory import MemorySaver
+        return MemorySaver()
diff --git a/app/core/tests/evaluation.py b/app/core/tests/evaluation.py
diff --git a/app/core/workflow/langraph_workflow.py b/app/core/workflow/langraph_workflow.py
diff --git a/app/data/big_benchmark.csv b/app/data/big_benchmark.csv
diff --git a/app/data/evaluation_dataset.csv b/app/data/evaluation_dataset.csv
diff --git a/app/tests/installation_test.py b/app/tests/installation_test.py
