truthstack_toolkit.py

"""
TruthStack LangChain Toolkit
Gives LangChain agents access to supplement-drug interaction checking.

Usage:
    from truthstack_toolkit import TruthStackToolkit

    toolkit = TruthStackToolkit(api_key="your-api-key")
    tools = toolkit.get_tools()
"""

from typing import Optional, Type
import requests
from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field


TRUTHSTACK_BASE_URL = "https://api.truthstack.co/api"


class TruthStackAPIClient:
    """HTTP client for TruthStack Vault API."""

    def __init__(self, api_key: str, base_url: str = TRUTHSTACK_BASE_URL):
        self.api_key = api_key
        self.base_url = base_url
        self.session = requests.Session()
        self.session.headers.update({
            "x-api-key": api_key,
            "Content-Type": "application/json",
        })

    def _get(self, endpoint: str, params: Optional[dict] = None) -> dict:
        url = f"{self.base_url}{endpoint}"
        resp = self.session.get(url, params=params, timeout=15)
        resp.raise_for_status()
        return resp.json()

    def _post(self, endpoint: str, data: dict) -> dict:
        url = f"{self.base_url}{endpoint}"
        resp = self.session.post(url, json=data, timeout=15)
        resp.raise_for_status()
        return resp.json()

    def search_compounds(self, query: str) -> dict:
        return self._get("/compounds/search", {"q": query})

    def get_compound(self, compound_id: str) -> dict:
        return self._get(f"/compounds/{compound_id}")

    def get_compound_interactions(self, compound_id: str) -> dict:
        return self._get(f"/compounds/{compound_id}/interactions")

    def check_stack(self, compounds: list[str]) -> dict:
        return self._post("/stack/check", {"compounds": compounds})

    def get_drug_profile(self, drug_name: str) -> dict:
        return self._get(f"/drugs/{drug_name}")

    def get_stats(self) -> dict:
        return self._get("/stats")


# --- Tool Input Schemas ---

class CheckInteractionsInput(BaseModel):
    compound: str = Field(description="Name of the supplement or compound to check (e.g., 'magnesium', 'ashwagandha', 'vitamin D')")

class SearchCompoundsInput(BaseModel):
    query: str = Field(description="Search query for compound name. Supports partial and fuzzy matching across 584 aliases.")

class GetCompoundInput(BaseModel):
    compound_id: str = Field(description="Compound ID or slug returned from search results")

class StackCheckInput(BaseModel):
    compounds: list[str] = Field(description="List of supplement/drug names to check for all pairwise interactions (e.g., ['magnesium', 'vitamin D', 'lisinopril'])")

class DrugProfileInput(BaseModel):
    drug_name: str = Field(description="Name of the drug to get interaction profile for (e.g., 'metformin', 'lisinopril', 'warfarin')")


# --- Tools ---

class CheckInteractionsTool(BaseTool):
    name: str = "truthstack_check_interactions"
    description: str = (
        "Check a supplement or compound for known interactions with drugs and other supplements. "
        "Returns interaction severity, evidence quality, and mechanism details. "
        "Use this when someone asks about the safety of a specific supplement."
    )
    args_schema: Type[BaseModel] = CheckInteractionsInput
    client: TruthStackAPIClient

    class Config:
        arbitrary_types_allowed = True

    def _run(self, compound: str) -> str:
        try:
            # First search for the compound
            search = self.client.search_compounds(compound)
            compounds = search.get("data", search.get("compounds", []))
            if not compounds:
                return f"No compound found matching '{compound}'. Try a different name or spelling."

            # Get the top match
            comp = compounds[0]
            comp_id = comp.get("id", comp.get("slug", comp.get("compound_id")))

            # Get interactions
            interactions = self.client.get_compound_interactions(comp_id)
            items = interactions.get("data", interactions.get("interactions", []))

            if not items:
                return f"'{comp.get('name', compound)}' found in database but has no known interactions recorded."

            # Format results
            lines = [f"Interactions for {comp.get('name', compound)} ({len(items)} found):"]
            for ix in items:
                severity = ix.get("severity", "unknown")
                interacts_with = ix.get("interacts_with", ix.get("compound_b", "unknown"))
                mechanism = ix.get("mechanism", ix.get("description", ""))
                lines.append(f"\n• {interacts_with} — Severity: {severity}")
                if mechanism:
                    lines.append(f"  Mechanism: {mechanism}")

            return "\n".join(lines)

        except requests.HTTPError as e:
            return f"API error checking interactions for '{compound}': {e}"
        except Exception as e:
            return f"Error: {e}"


class SearchCompoundsTool(BaseTool):
    name: str = "truthstack_search_compounds"
    description: str = (
        "Search the TruthStack database for a supplement or compound by name. "
        "Supports fuzzy matching across 584 aliases. Returns compound IDs, names, and categories."
    )
    args_schema: Type[BaseModel] = SearchCompoundsInput
    client: TruthStackAPIClient

    class Config:
        arbitrary_types_allowed = True

    def _run(self, query: str) -> str:
        try:
            results = self.client.search_compounds(query)
            compounds = results.get("data", results.get("compounds", []))
            if not compounds:
                return f"No compounds found matching '{query}'."

            lines = [f"Found {len(compounds)} compounds matching '{query}':"]
            for c in compounds[:10]:
                name = c.get("name", "unknown")
                comp_id = c.get("id", c.get("slug", ""))
                category = c.get("category", "")
                aliases = c.get("aliases", [])
                line = f"• {name} (ID: {comp_id})"
                if category:
                    line += f" — {category}"
                if aliases:
                    line += f" — also: {', '.join(aliases[:3])}"
                lines.append(line)

            return "\n".join(lines)

        except Exception as e:
            return f"Error searching compounds: {e}"


class GetCompoundTool(BaseTool):
    name: str = "truthstack_get_compound"
    description: str = (
        "Get the full profile of a compound including all known interactions, evidence, "
        "and safety information. Use after searching to get detailed information."
    )
    args_schema: Type[BaseModel] = GetCompoundInput
    client: TruthStackAPIClient

    class Config:
        arbitrary_types_allowed = True

    def _run(self, compound_id: str) -> str:
        try:
            result = self.client.get_compound(compound_id)
            compound = result.get("data", result)
            name = compound.get("name", compound_id)

            lines = [f"Compound: {name}"]
            if compound.get("category"):
                lines.append(f"Category: {compound['category']}")
            if compound.get("description"):
                lines.append(f"Description: {compound['description']}")
            if compound.get("aliases"):
                lines.append(f"Also known as: {', '.join(compound['aliases'])}")

            interactions = compound.get("interactions", [])
            if interactions:
                lines.append(f"\nKnown interactions ({len(interactions)}):")
                for ix in interactions:
                    severity = ix.get("severity", "unknown")
                    target = ix.get("interacts_with", ix.get("compound_b", "unknown"))
                    lines.append(f"  • {target} — {severity}")

            return "\n".join(lines)

        except requests.HTTPError as e:
            if e.response and e.response.status_code == 404:
                return f"Compound '{compound_id}' not found."
            return f"API error: {e}"
        except Exception as e:
            return f"Error: {e}"


class StackCheckTool(BaseTool):
    name: str = "truthstack_stack_check"
    description: str = (
        "Check an entire supplement stack for all pairwise interactions. "
        "Pass a list of supplement and drug names. Returns all flagged interactions between them. "
        "Use this when someone shares their full supplement regimen or asks about combining multiple supplements."
    )
    args_schema: Type[BaseModel] = StackCheckInput
    client: TruthStackAPIClient

    class Config:
        arbitrary_types_allowed = True

    def _run(self, compounds: list[str]) -> str:
        try:
            result = self.client.check_stack(compounds)
            interactions = result.get("data", result.get("interactions", []))
            flagged = [i for i in interactions if i.get("severity") != "none"]

            if not flagged:
                return f"No interactions found between: {', '.join(compounds)}. Stack appears safe based on current data."

            lines = [f"⚠️ Found {len(flagged)} interactions in stack ({', '.join(compounds)}):"]
            for ix in flagged:
                a = ix.get("compound_a", ix.get("from", "?"))
                b = ix.get("compound_b", ix.get("to", "?"))
                severity = ix.get("severity", "unknown")
                mechanism = ix.get("mechanism", ix.get("description", ""))
                lines.append(f"\n• {a} ↔ {b} — Severity: {severity}")
                if mechanism:
                    lines.append(f"  {mechanism}")

            return "\n".join(lines)

        except Exception as e:
            return f"Error checking stack: {e}"


class DrugProfileTool(BaseTool):
    name: str = "truthstack_get_drug_profile"
    description: str = (
        "Get a drug's interaction profile showing all known supplement interactions. "
        "Covers 25 common drugs including statins, SSRIs, blood thinners, and metformin."
    )
    args_schema: Type[BaseModel] = DrugProfileInput
    client: TruthStackAPIClient

    class Config:
        arbitrary_types_allowed = True

    def _run(self, drug_name: str) -> str:
        try:
            result = self.client.get_drug_profile(drug_name)
            drug = result.get("data", result)
            name = drug.get("name", drug_name)

            lines = [f"Drug Profile: {name}"]
            if drug.get("drug_class"):
                lines.append(f"Class: {drug['drug_class']}")
            if drug.get("common_uses"):
                lines.append(f"Common uses: {drug['common_uses']}")

            interactions = drug.get("interactions", [])
            if interactions:
                lines.append(f"\nSupplement interactions ({len(interactions)}):")
                for ix in interactions:
                    supp = ix.get("supplement", ix.get("compound", "unknown"))
                    severity = ix.get("severity", "unknown")
                    lines.append(f"  • {supp} — {severity}")

            warnings = drug.get("warnings", [])
            if warnings:
                lines.append(f"\nWarnings:")
                for w in warnings:
                    lines.append(f"  ⚠️ {w}")

            return "\n".join(lines)

        except requests.HTTPError as e:
            if e.response and e.response.status_code == 404:
                return f"No drug profile found for '{drug_name}'. Currently covering 25 common drugs."
            return f"API error: {e}"
        except Exception as e:
            return f"Error: {e}"


# --- Toolkit ---

class TruthStackToolkit:
    """
    LangChain toolkit for TruthStack supplement-drug interaction API.

    Provides 5 tools for checking supplement safety, searching compounds,
    and querying drug interaction profiles.

    Usage:
        toolkit = TruthStackToolkit(api_key="your-api-key")
        tools = toolkit.get_tools()
    """

    def __init__(self, api_key: str, base_url: str = TRUTHSTACK_BASE_URL):
        self.client = TruthStackAPIClient(api_key=api_key, base_url=base_url)

    def get_tools(self) -> list[BaseTool]:
        return [
            CheckInteractionsTool(client=self.client),
            SearchCompoundsTool(client=self.client),
            GetCompoundTool(client=self.client),
            StackCheckTool(client=self.client),
            DrugProfileTool(client=self.client),
        ]