initial commit

venv/Lib/site-packages/langchain_community/utilities/dataherald.py

@@ -0,0 +1,68 @@
+"""Util that calls Dataherald."""
+
+from typing import Any, Dict, Optional
+
+from langchain_core.utils import get_from_dict_or_env
+from pydantic import BaseModel, ConfigDict, model_validator
+
+
+class DataheraldAPIWrapper(BaseModel):
+    """Wrapper for Dataherald.
+
+    Docs for using:
+
+    1. Go to dataherald and sign up
+    2. Create an API key
+    3. Save your API key into DATAHERALD_API_KEY env variable
+    4. pip install dataherald
+
+    """
+
+    dataherald_client: Any = None  #: :meta private:
+    db_connection_id: str
+    dataherald_api_key: Optional[str] = None
+
+    model_config = ConfigDict(
+        extra="forbid",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_environment(cls, values: Dict) -> Any:
+        """Validate that api key and python package exists in environment."""
+        dataherald_api_key = get_from_dict_or_env(
+            values, "dataherald_api_key", "DATAHERALD_API_KEY"
+        )
+        values["dataherald_api_key"] = dataherald_api_key
+
+        try:
+            import dataherald
+
+        except ImportError:
+            raise ImportError(
+                "dataherald is not installed. "
+                "Please install it with `pip install dataherald`"
+            )
+
+        client = dataherald.Dataherald(api_key=dataherald_api_key)
+        values["dataherald_client"] = client
+
+        return values
+
+    def run(self, prompt: str) -> str:
+        """Generate a sql query through Dataherald and parse result."""
+        from dataherald.types.sql_generation_create_params import Prompt
+
+        prompt_obj = Prompt(text=prompt, db_connection_id=self.db_connection_id)
+        res = self.dataherald_client.sql_generations.create(prompt=prompt_obj)
+
+        try:
+            answer = res.sql
+            if not answer:
+                # We don't want to return the assumption alone if answer is empty
+                return "No answer"
+            else:
+                return f"Answer: {answer}"
+
+        except StopIteration:
+            return "Dataherald wasn't able to answer it"