Upload agent/toxicologist_agent.py with huggingface_hub

agent/toxicologist_agent.py

@@ -0,0 +1,85 @@
+import json
+import os
+from groq import Groq
+from pydantic import BaseModel, Field, TypeAdapter
+from typing import List, Optional
+from config import GROQ_API_KEY, GROQ_MODEL
+
+
+class ToxicologyReport(BaseModel):
+    target_assay: str = Field(description="The target assay endpoint evaluated.")
+    risk_level: str = Field(description="Overall risk level: Low, Moderate, High, or Critical.")
+    toxicophore_assessment: str = Field(description="Analysis of which structural features contribute to predicted toxicity.")
+    biochemical_mechanism: str = Field(description="Explanation of the biochemical mechanism linking structure to endpoint.")
+    structural_alerts_found: List[dict] = Field(description="List of structural alerts matched and their relevance.")
+    bioisosteric_replacements: List[str] = Field(description="Suggested safer bioisosteric replacements.")
+    confidence: str = Field(description="Confidence in this assessment based on model prediction and alert matching.")
+
+
+class ToxicologistAgent:
+    def __init__(self):
+        self.client = Groq(api_key=GROQ_API_KEY)
+        self.model = GROQ_MODEL
+        self.schema_adapter = TypeAdapter(ToxicologyReport)
+        self.json_schema = self._build_strict_schema()
+
+    def _build_strict_schema(self) -> dict:
+        raw = self.schema_adapter.json_schema()
+        raw["additionalProperties"] = False
+        if "properties" in raw:
+            for prop_val in raw["properties"].values():
+                if isinstance(prop_val, dict) and "properties" in prop_val:
+                    prop_val["additionalProperties"] = False
+        if "$defs" in raw:
+            for def_val in raw["$defs"].values():
+                if isinstance(def_val, dict):
+                    def_val["additionalProperties"] = False
+        return raw
+
+    def generate_report(self, smiles: str, target_assay: str, predictions: list,
+                        shap_attributions: list, high_attr_indices: list,
+                        alerts_found: list) -> ToxicologyReport:
+        prompt = f"""You are an expert computational toxicologist. Analyze the following compound's toxicological profile.
+
+COMPOUND SMILES: {smiles}
+TARGET ASSAY: {target_assay} ({predictions[0].get('target_class', 'Unknown')})
+PREDICTED PROBABILITY: {predictions[0].get('probability', 0.5):.4f}
+PREDICTED CLASS: {predictions[0].get('predicted_class', 'Unknown')}
+
+SHAP HIGH-ATTRIBUTION ATOM INDICES: {high_attr_indices}
+FULL SHAP VECTOR: {shap_attributions}
+
+STRUCTURAL ALERTS DETECTED: {json.dumps(alerts_found, indent=2)}
+
+Based on this data:
+1. Identify which specific structural fragments (toxicophores) are driving the {predictions[0].get('predicted_class', 'Unknown')} prediction for {target_assay}.
+2. Explain the biochemical mechanism linking these fragments to the assay endpoint.
+3. Suggest 2-3 specific bioisosteric replacements that could mitigate toxicity while preserving core scaffold.
+4. Assess overall risk level.
+5. Provide confidence in the assessment.
+
+Output a structured report adhering to the provided schema.
+"""
+        chat_completion = self.client.chat.completions.create(
+            model=self.model,
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a professional computational toxicologist. Output reports strictly adhering to the provided JSON schema."
+                },
+                {"role": "user", "content": prompt}
+            ],
+            response_format={
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "Toxicology_Report",
+                    "strict": True,
+                    "schema": self.json_schema
+                }
+            },
+            temperature=0.1,
+        )
+
+        raw_json = chat_completion.choices[0].message.content
+        report = self.schema_adapter.validate_json(raw_json)
+        return report