Update tools/umls_tool.py

tools/umls_tool.py

@@ -0,0 +1,36 @@
+from langchain.tools import BaseTool
+from typing import Type
+from pydantic import BaseModel, Field
+from clinical_nlp.umls_bioportal import search_umls_term
+from services.logger import app_logger
+from services.metrics import log_tool_usage
+
+class UMLSInput(BaseModel):
+    term: str = Field(description="The medical term to search for in UMLS.")
+
+class UMLSLookupTool(BaseTool):
+    name: str = "umls_lookup"
+    description: str = (
+        "Useful for looking up medical terms, concepts, and their definitions or CUIs "
+        "(Concept Unique Identifiers) in the Unified Medical Language System (UMLS). "
+        "Provide a specific medical term."
+    )
+    args_schema: Type[BaseModel] = UMLSInput
+
+    def _run(self, term: str) -> str:
+        app_logger.info(f"UMLS Tool called with term: {term}")
+        log_tool_usage(self.name)
+        results = search_umls_term(term)
+        if "error" in results:
+            return f"Error from UMLS lookup: {results['error']}"
+        # Format results for LLM consumption
+        if results.get("results"):
+            formatted_results = []
+            for res in results["results"][:3]: # Limit to 3 results
+                 formatted_results.append(f"- Term: {res.get('name', 'N/A')}, CUI: {res.get('ui', 'N/A')}") # Assuming 'ui' for CUI based on UMLS API
+            return "UMLS Results:\n" + "\n".join(formatted_results) if formatted_results else "No results found in UMLS."
+        return "No results found or unexpected format from UMLS."
+
+    async def _arun(self, term: str) -> str:
+        # For simplicity, using sync version for now
+        return self._run(term)