Update tools/umls_tool.py

tools/umls_tool.py

@@ -1,36 +1,62 @@
-from langchain.tools import BaseTool
-from typing import Type
+# /home/user/app/tools/umls_tool.py
+from langchain_core.tools import BaseTool
+from typing import Type, Optional # Added Optional
 from pydantic import BaseModel, Field
+# Assuming clinical_nlp.umls_bioportal.search_umls_term exists and works
 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.")
+    term: str = Field(description="The specific medical term to search for in UMLS.")
+    # Example: search_type: Optional[str] = Field(default="exact", description="Type of search, e.g., 'exact', 'approximate'.")
 
 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."
+        "Use this tool to find definitions, Concept Unique Identifiers (CUIs), and related concepts for specific medical terms "
+        "from the Unified Medical Language System (UMLS). Best for precise terminology lookup. "
+        "Input should be a single medical term (e.g., 'myocardial infarction')."
     )
     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"):
+    def _run(self, term: str) -> str: # Add other args from UMLSInput if defined
+        app_logger.info(f"UMLS Tool called with term: '{term}'")
+        log_tool_usage(self.name, {"query_term": term}) # Log term for metrics
+
+        if not term or not term.strip():
+            return "Error from UMLS lookup: No search term provided."
+
+        try:
+            # Pass additional args like search_type to search_umls_term if your function supports them
+            results = search_umls_term(term) # Example: search_umls_term(term, search_type=search_type)
+        except Exception as e:
+            app_logger.error(f"Exception during UMLS search for '{term}': {e}", exc_info=True)
+            return f"Error performing UMLS lookup for '{term}': An unexpected error occurred during the search."
+
+        if isinstance(results, dict) and "error" in results:
+            app_logger.warning(f"UMLS lookup for '{term}' returned an error: {results['error']}")
+            return f"Error from UMLS lookup for '{term}': {results['error']}"
+        
+        if isinstance(results, dict) and 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."
+            for res_item in results["results"][:3]: # Limit to top 3 results for brevity
+                name = res_item.get('name', 'N/A')
+                cui = res_item.get('ui', 'N/A') # Assuming 'ui' is the CUI field from your API response
+                # You might want to add definition or semantic type if available
+                # definition = res_item.get('definition', 'No definition available.')
+                formatted_results.append(f"- Term: {name}, CUI: {cui}")
+            
+            if formatted_results:
+                return f"UMLS Search Results for '{term}':\n" + "\n".join(formatted_results)
+            else:
+                return f"No specific items found in UMLS results for '{term}', though the query was successful."
+        
+        app_logger.warning(f"No results or unexpected format from UMLS for '{term}'. Raw: {str(results)[:200]}")
+        return f"No distinct results found or unexpected response format from UMLS for '{term}'."
 
     async def _arun(self, term: str) -> str:
-        # For simplicity, using sync version for now
+        # For simplicity, using sync version. Implement true async if search_umls_term is async.
+        # import asyncio
+        # return await asyncio.to_thread(self._run, term)
         return self._run(term)