khora 0.0.1__py3-none-any.whl

khora/tools/api_tool.py

@@ -0,0 +1,81 @@
+"""API tool for fetching data from REST APIs."""
+
+import json
+from typing import Any, Dict, Optional
+
+import httpx
+from langchain.tools import BaseTool
+from pydantic import Field
+
+
+class APITool(BaseTool):
+    """Tool for making API requests based on AI-generated specifications."""
+
+    name: str = "api_fetcher"
+    description: str = (
+        "Fetch data from APIs. The tool accepts a URL, HTTP method, "
+        "headers, and optional body/params based on the AI prompt analysis."
+    )
+
+    timeout: int = Field(default=30, description="Request timeout in seconds")
+
+    def _run(
+        self,
+        url: str,
+        method: str = "GET",
+        headers: Optional[Dict[str, str]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        json_body: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Execute API request.
+
+        Args:
+            url: API endpoint URL
+            method: HTTP method (GET, POST, PUT, DELETE, etc.)
+            headers: Optional HTTP headers
+            params: Optional query parameters
+            json_body: Optional JSON body for POST/PUT requests
+
+        Returns:
+            Response data as dictionary
+        """
+        try:
+            with httpx.Client(timeout=self.timeout) as client:
+                response = client.request(
+                    method=method.upper(),
+                    url=url,
+                    headers=headers,
+                    params=params,
+                    json=json_body,
+                )
+
+                response.raise_for_status()
+
+                # Try to parse JSON response
+                try:
+                    data = response.json()
+                except json.JSONDecodeError:
+                    data = {"text": response.text}
+
+                return {
+                    "status": "success",
+                    "status_code": response.status_code,
+                    "data": data,
+                    "headers": dict(response.headers),
+                }
+
+        except httpx.HTTPStatusError as e:
+            return {
+                "status": "error",
+                "status_code": e.response.status_code,
+                "error": str(e),
+                "response_text": e.response.text,
+            }
+        except Exception as e:
+            return {"status": "error", "error": str(e), "error_type": type(e).__name__}
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+        """Async version of the API tool."""
+        raise NotImplementedError("Async execution not implemented yet")

khora/tools/google_docs_tool.py

@@ -0,0 +1,169 @@
+"""Google Docs and Sheets tool for extracting data."""
+
+from typing import Any, Dict, List, Optional
+
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from langchain.tools import BaseTool
+from pydantic import Field
+
+
+class GoogleDocsTool(BaseTool):
+    """Tool for extracting data from Google Docs and Sheets."""
+
+    name: str = "google_docs_fetcher"
+    description: str = (
+        "Extract data from Google Docs and Google Sheets. "
+        "Requires document/sheet ID and appropriate permissions."
+    )
+
+    credentials_path: Optional[str] = Field(
+        default=None, description="Path to Google service account credentials JSON"
+    )
+    scopes: List[str] = Field(
+        default_factory=lambda: [
+            "https://www.googleapis.com/auth/documents.readonly",
+            "https://www.googleapis.com/auth/spreadsheets.readonly",
+            "https://www.googleapis.com/auth/drive.readonly",
+        ]
+    )
+
+    def _run(
+        self,
+        document_id: str,
+        document_type: str = "sheet",
+        sheet_range: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Extract data from Google Docs or Sheets.
+
+        Args:
+            document_id: Google document or sheet ID
+            document_type: Type of document ("doc" or "sheet")
+            sheet_range: For sheets, the A1 notation range (e.g., "Sheet1!A1:D10")
+
+        Returns:
+            Extracted data as dictionary
+        """
+        try:
+            # Initialize credentials
+            if self.credentials_path:
+                credentials = service_account.Credentials.from_service_account_file(  # type: ignore
+                    self.credentials_path, scopes=self.scopes
+                )
+            else:
+                # Use default credentials if available
+                credentials = None
+
+            if document_type.lower() == "sheet":
+                return self._extract_sheet_data(document_id, sheet_range, credentials)
+            elif document_type.lower() == "doc":
+                return self._extract_doc_data(document_id, credentials)
+            else:
+                return {
+                    "status": "error",
+                    "error": f"Unsupported document type: {document_type}",
+                }
+
+        except Exception as e:
+            return {"status": "error", "error": str(e), "error_type": type(e).__name__}
+
+    def _extract_sheet_data(
+        self, sheet_id: str, sheet_range: Optional[str], credentials: Any
+    ) -> Dict[str, Any]:
+        """Extract data from Google Sheets."""
+        service = build("sheets", "v4", credentials=credentials)
+
+        # Get sheet metadata
+        sheet_metadata = service.spreadsheets().get(spreadsheetId=sheet_id).execute()
+
+        sheets = sheet_metadata.get("sheets", [])
+
+        result = {
+            "title": sheet_metadata.get("properties", {}).get("title"),
+            "sheets": [s["properties"]["title"] for s in sheets],
+            "data": {},
+        }
+
+        # If no range specified, get all sheets
+        if not sheet_range:
+            for sheet in sheets:
+                sheet_name = sheet["properties"]["title"]
+                range_name = f"{sheet_name}!A:Z"
+                try:
+                    sheet_data = (
+                        service.spreadsheets()
+                        .values()
+                        .get(spreadsheetId=sheet_id, range=range_name)
+                        .execute()
+                    )
+
+                    values = sheet_data.get("values", [])
+                    if values:
+                        # Convert to list of dicts using first row as headers
+                        headers = values[0] if values else []
+                        rows = []
+                        for row in values[1:]:
+                            row_dict = {}
+                            for i, header in enumerate(headers):
+                                row_dict[header] = row[i] if i < len(row) else ""
+                            rows.append(row_dict)
+                        result["data"][sheet_name] = rows
+                except Exception:
+                    result["data"][sheet_name] = []
+        else:
+            # Get specific range
+            sheet_data = (
+                service.spreadsheets()
+                .values()
+                .get(spreadsheetId=sheet_id, range=sheet_range)
+                .execute()
+            )
+
+            values = sheet_data.get("values", [])
+            if values:
+                headers = values[0] if values else []
+                rows = []
+                for row in values[1:]:
+                    row_dict = {}
+                    for i, header in enumerate(headers):
+                        row_dict[header] = row[i] if i < len(row) else ""
+                    rows.append(row_dict)
+                result["data"]["requested_range"] = rows
+
+        return {"status": "success", "data": result}
+
+    def _extract_doc_data(self, doc_id: str, credentials: Any) -> Dict[str, Any]:
+        """Extract data from Google Docs."""
+        service = build("docs", "v1", credentials=credentials)
+
+        # Get document
+        document = service.documents().get(documentId=doc_id).execute()
+
+        title = document.get("title")
+        content = []
+
+        # Extract text content
+        for element in document.get("body", {}).get("content", []):
+            if "paragraph" in element:
+                paragraph = element["paragraph"]
+                text_elements = []
+                for elem in paragraph.get("elements", []):
+                    if "textRun" in elem:
+                        text_elements.append(elem["textRun"]["content"])
+                if text_elements:
+                    content.append("".join(text_elements))
+
+        return {
+            "status": "success",
+            "data": {
+                "title": title,
+                "content": "\n".join(content),
+                "document_id": doc_id,
+            },
+        }
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+        """Async version of the Google Docs tool."""
+        raise NotImplementedError("Async execution not implemented yet")

khora/tools/web_scraper_tool.py

@@ -0,0 +1,197 @@
+"""Web scraper tool for extracting data from websites using Playwright."""
+
+import asyncio
+from typing import Any, Dict, Optional
+
+from langchain.tools import BaseTool
+from playwright.async_api import async_playwright
+from pydantic import Field
+
+
+class WebScraperTool(BaseTool):
+    """Tool for scraping data from websites using Playwright."""
+
+    name: str = "web_scraper"
+    description: str = (
+        "Extract data from websites using Playwright. Can handle JavaScript-rendered "
+        "content, interact with pages, and extract complex data structures."
+    )
+
+    timeout: int = Field(default=30000, description="Page timeout in milliseconds")
+    headless: bool = Field(default=True, description="Run browser in headless mode")
+
+    def _run(
+        self,
+        url: str,
+        wait_for: Optional[str] = None,
+        selectors: Optional[Dict[str, str]] = None,
+        extract_all_text: bool = False,
+        extract_links: bool = False,
+        extract_tables: bool = False,
+        screenshot: bool = False,
+        execute_script: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Scrape web page and extract data using Playwright.
+
+        Args:
+            url: Website URL to scrape
+            wait_for: CSS selector or state to wait for before extraction
+            selectors: CSS selectors for specific elements
+            extract_all_text: Extract all text content
+            extract_links: Extract all links
+            extract_tables: Extract tables as structured data
+            screenshot: Take a screenshot of the page
+            execute_script: Custom JavaScript to execute on the page
+
+        Returns:
+            Extracted data as dictionary
+        """
+        # Run async function in sync context
+        return asyncio.run(
+            self._async_run(
+                url=url,
+                wait_for=wait_for,
+                selectors=selectors,
+                extract_all_text=extract_all_text,
+                extract_links=extract_links,
+                extract_tables=extract_tables,
+                screenshot=screenshot,
+                execute_script=execute_script,
+                **kwargs,
+            )
+        )
+
+    async def _async_run(
+        self,
+        url: str,
+        wait_for: Optional[str] = None,
+        selectors: Optional[Dict[str, str]] = None,
+        extract_all_text: bool = False,
+        extract_links: bool = False,
+        extract_tables: bool = False,
+        screenshot: bool = False,
+        execute_script: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """Async implementation of web scraping."""
+        try:
+            async with async_playwright() as p:
+                # Launch browser
+                browser = await p.chromium.launch(headless=self.headless)
+                context = await browser.new_context(
+                    viewport={"width": 1920, "height": 1080},
+                    user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
+                )
+                page = await context.new_page()
+
+                # Navigate to URL
+                await page.goto(url, timeout=self.timeout)
+
+                # Wait for specific element or state if specified
+                if wait_for:
+                    if wait_for in ["load", "domcontentloaded", "networkidle"]:
+                        await page.wait_for_load_state(wait_for)  # type: ignore
+                    else:
+                        await page.wait_for_selector(wait_for, timeout=self.timeout)
+                else:
+                    # Default: wait for network to be idle
+                    await page.wait_for_load_state("networkidle", timeout=self.timeout)
+
+                result: Dict[str, Any] = {"url": page.url, "title": await page.title()}
+
+                # Execute custom JavaScript if provided
+                if execute_script:
+                    script_result = await page.evaluate(execute_script)
+                    result["script_result"] = script_result
+
+                # Extract based on selectors
+                if selectors:
+                    extracted_data: Dict[str, list[str]] = {}
+                    for key, selector in selectors.items():
+                        elements = await page.query_selector_all(selector)
+                        extracted_data[key] = []
+                        for elem in elements:
+                            text = await elem.text_content()
+                            if text:
+                                extracted_data[key].append(text.strip())
+                    result["selected_data"] = extracted_data
+
+                # Extract all text
+                if extract_all_text:
+                    result["text"] = await page.inner_text("body")
+
+                # Extract links
+                if extract_links:
+                    links = await page.evaluate(
+                        """
+                        () => Array.from(document.querySelectorAll('a[href]')).map(a => ({
+                            text: a.textContent.trim(),
+                            href: a.href,
+                            title: a.title || null
+                        }))
+                    """
+                    )
+                    result["links"] = links
+
+                # Extract tables
+                if extract_tables:
+                    tables = await page.evaluate(
+                        """
+                        () => Array.from(document.querySelectorAll('table')).map(table => {
+                            const headers = Array.from(table.querySelectorAll('th')).map(th => th.textContent.trim());
+                            const rows = Array.from(table.querySelectorAll('tr')).slice(1).map(row => {
+                                const cells = Array.from(row.querySelectorAll('td, th'));
+                                const rowData = {};
+                                cells.forEach((cell, i) => {
+                                    const key = headers[i] || `column_${i}`;
+                                    rowData[key] = cell.textContent.trim();
+                                });
+                                return rowData;
+                            });
+                            return rows;
+                        })
+                    """
+                    )
+                    result["tables"] = tables
+
+                # Take screenshot if requested
+                if screenshot:
+                    screenshot_data = await page.screenshot(full_page=True)
+                    result["screenshot"] = {
+                        "size": len(screenshot_data),
+                        "note": "Screenshot data available as bytes",
+                    }
+
+                await browser.close()
+
+                return {"status": "success", "data": result}
+
+        except Exception as e:
+            return {"status": "error", "error": str(e), "error_type": type(e).__name__}
+
+    async def _arun(
+        self,
+        url: str,
+        wait_for: Optional[str] = None,
+        selectors: Optional[Dict[str, str]] = None,
+        extract_all_text: bool = False,
+        extract_links: bool = False,
+        extract_tables: bool = False,
+        screenshot: bool = False,
+        execute_script: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """Async version of the web scraper tool."""
+        return await self._async_run(
+            url=url,
+            wait_for=wait_for,
+            selectors=selectors,
+            extract_all_text=extract_all_text,
+            extract_links=extract_links,
+            extract_tables=extract_tables,
+            screenshot=screenshot,
+            execute_script=execute_script,
+            **kwargs,
+        )

khora/utils/__init__.py

@@ -0,0 +1,6 @@
+"""Utility functions and helpers."""
+
+from .config import load_config
+from .data_models import DataRequest, DataResponse
+
+__all__ = ["load_config", "DataRequest", "DataResponse"]

khora/utils/config.py

@@ -0,0 +1,54 @@
+"""Configuration management for Khora."""
+
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from dotenv import load_dotenv
+
+
+def load_config(config_path: Optional[Path] = None) -> Dict[str, Any]:
+    """
+    Load configuration from environment variables and optional config file.
+
+    Args:
+        config_path: Optional path to configuration file
+
+    Returns:
+        Configuration dictionary
+    """
+    # Load environment variables
+    load_dotenv()
+
+    config = {
+        # OpenAI Configuration
+        "openai_api_key": os.getenv("OPENAI_API_KEY"),
+        "openai_model": os.getenv("OPENAI_MODEL", "gpt-4-turbo-preview"),
+        # Google Configuration
+        "google_credentials_path": os.getenv("GOOGLE_CREDENTIALS_PATH"),
+        "google_scopes": [
+            "https://www.googleapis.com/auth/documents.readonly",
+            "https://www.googleapis.com/auth/spreadsheets.readonly",
+            "https://www.googleapis.com/auth/drive.readonly",
+        ],
+        # Dagster Configuration
+        "dagster_home": os.getenv("DAGSTER_HOME", "/tmp/dagster"),
+        "dagster_storage": {
+            "postgres": {
+                "postgres_db": os.getenv("DAGSTER_PG_DB", "dagster"),
+                "postgres_host": os.getenv("DAGSTER_PG_HOST", "localhost"),
+                "postgres_port": int(os.getenv("DAGSTER_PG_PORT", "5432")),
+                "postgres_user": os.getenv("DAGSTER_PG_USER", "dagster"),
+                "postgres_password": os.getenv("DAGSTER_PG_PASSWORD", ""),
+            }
+        },
+        # General Configuration
+        "log_level": os.getenv("LOG_LEVEL", "INFO"),
+        "cache_enabled": os.getenv("CACHE_ENABLED", "true").lower() == "true",
+        "cache_ttl": int(os.getenv("CACHE_TTL", "3600")),
+    }
+
+    # Remove None values
+    config = {k: v for k, v in config.items() if v is not None}
+
+    return config

khora/utils/data_models.py

@@ -0,0 +1,57 @@
+"""Data models for Khora pipeline operations."""
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class DataSourceType(str, Enum):
+    """Supported data source types."""
+
+    API = "api"
+    WEB_SCRAPER = "web_scraper"
+    GOOGLE_DOCS = "google_docs"
+    SPREADSHEET = "spreadsheet"
+
+
+class DataRequest(BaseModel):
+    """Model for data fetching requests."""
+
+    source_type: DataSourceType
+    prompt: str = Field(..., description="AI prompt describing what data to fetch")
+    source_config: Dict[str, Any] = Field(
+        default_factory=dict, description="Configuration specific to the data source"
+    )
+    filters: Optional[Dict[str, Any]] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+    model_config = ConfigDict(use_enum_values=True)
+
+
+class DataResponse(BaseModel):
+    """Model for data fetching responses."""
+
+    request_id: str
+    status: str = Field(..., description="success, error, or partial")
+    data: Optional[Any] = None
+    error_message: Optional[str] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    source_type: DataSourceType
+
+    model_config = ConfigDict(use_enum_values=True)
+
+
+class PipelineConfig(BaseModel):
+    """Configuration for pipeline execution."""
+
+    name: str
+    description: Optional[str] = None
+    requests: List[DataRequest]
+    parallel_execution: bool = True
+    retry_config: Dict[str, Any] = Field(
+        default_factory=lambda: {"max_retries": 3, "retry_delay": 5}
+    )
+    output_format: str = "json"