simplai-sdk 0.1.0__py3-none-any.whl

simplai_sdk/simplai.py

@@ -0,0 +1,239 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from core import (
+    cancel_bulk_run,
+    cancel_execution,
+    download_bulk_run_result,
+    execute_and_wait_workflow,
+    execute_workflow,
+    get_bulk_run_status,
+    get_tool_result,
+    trigger_bulk_run,
+)
+from core.agents import (
+    agent_chat_async,
+    agent_chat_stream_async,
+    AgentResult,
+    AgentStreamChunk,
+)
+from core.workflows.scheduling import cancel_schedule_run as cancel_schedule_run_func, schedule_run as schedule_run_func
+from billing import get_billing_detail
+from traces.agents import get_trace_tree, TraceNode
+from utils.config import get_api_key, get_base_url
+
+
+class SimplAI:
+    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
+        """
+        Initialize SimplAI client.
+
+        Args:
+            api_key: Optional API key. If not provided, loads from API_KEY env var or .env file.
+            base_url: Optional base URL. If not provided, loads from BASE_URL env var or uses default.
+
+        Raises:
+            ValueError: If API key is not found in any source.
+        """
+        self.api_key = get_api_key(api_key)
+
+    # ------------------------------------------------------------------
+    # WORKFLOWS
+    # ------------------------------------------------------------------
+
+    async def execute_workflow(
+        self,
+        workflow_id: str,
+        inputs: Dict[str, Any],
+    ) -> Dict[str, str]:
+        return await execute_workflow(workflow_id, inputs, self.api_key)
+
+    async def trigger_bulk_run(
+        self,
+        workflow_id: str,
+        file_link: str,
+        batch_size: int,
+        input_mapping: List[Dict[str, str]],
+    ) -> Dict[str, str]:
+        return await trigger_bulk_run(
+            workflow_id,
+            file_link,
+            batch_size,
+            input_mapping,
+            self.api_key,
+        )
+
+    async def execute_and_wait_workflow(
+        self,
+        workflow_id: str,
+        inputs: Dict[str, Any],
+        timeout: float = 60.0,
+        poll_interval: float = 2.0,
+    ) -> Dict[str, Any]:
+        return await execute_and_wait_workflow(
+            workflow_id,
+            inputs,
+            self.api_key,
+            timeout,
+            poll_interval,
+        )
+
+    async def get_tool_result(self, execution_id: str) -> Dict[str, Any]:
+        return await get_tool_result(execution_id, self.api_key)
+
+    async def cancel_execution(self, execution_id: str) -> Dict[str, str]:
+        return await cancel_execution(execution_id, self.api_key)
+
+    async def get_bulk_run_status(self, bulk_run_id: str) -> Dict[str, str]:
+        return await get_bulk_run_status(bulk_run_id, self.api_key)
+
+    async def cancel_bulk_run(self, bulk_run_id: str) -> Dict[str, str]:
+        return await cancel_bulk_run(bulk_run_id, self.api_key)
+
+    async def download_bulk_run_result(self, bulk_run_id: str) -> str:
+        """
+        Download the bulk run result as CSV text.
+        """
+        return await download_bulk_run_result(bulk_run_id, self.api_key)
+
+  
+
+    async def agent_chat(
+        self,
+        agent_id: str,
+        message: str,
+    ) -> AgentResult:
+        """
+        High-level agent chat helper (async), matching workflow style.
+
+        Only required params here (agent_id, message).
+        All optional config is handled by core.agents via defaults.
+        """
+        return await agent_chat_async(
+            agent_id=agent_id,
+            message=message,
+            api_key=self.api_key,
+        )
+
+    async def agent_chat_stream(
+        self,
+        agent_id: str,
+        message: str,
+        on_chunk: Optional[Callable[[AgentStreamChunk], None]] = None,
+    ) -> AgentResult:
+        """
+        High-level streaming agent chat helper (async), workflow-style.
+
+        Only required call-specific params here.
+        All optional config is handled by core.agents via defaults.
+        
+        Note: The first chunk will contain trace_id, node_id, and tree_id for tracing.
+        Access them via chunk.trace_id, chunk.node_id, and chunk.tree_id in the on_chunk callback.
+        Use get_trace_tree to fetch trace details.
+        """
+        return await agent_chat_stream_async(
+            agent_id=agent_id,
+            message=message,
+            api_key=self.api_key,
+            on_chunk=on_chunk,
+        )
+
+    # ------------------------------------------------------------------
+    # TRACING
+    # ------------------------------------------------------------------
+
+    def get_trace_tree(
+        self,
+        node_id: str,
+        trace_id: Optional[str] = None,
+        tree_id: Optional[str] = None,
+        max_depth: Optional[int] = None,
+        user_id: Optional[str] = None,
+        tenant_id: str = "1",
+        project_id: Optional[int] = None,
+        seller_id: Optional[str] = None,
+        client_id: Optional[str] = None,
+        seller_profile_id: Optional[str] = None,
+        raw_response: bool = False,
+    ) -> Union[TraceNode, Dict[str, Any]]:
+        """Get a trace sub-tree starting from a specific node.
+
+        Args:
+            node_id: Node ID to start the tree from.
+            trace_id: Trace ID (same as tree_id). Either trace_id or tree_id must be provided.
+            tree_id: Tree ID of the trace (same as trace_id). Either trace_id or tree_id must be provided.
+            max_depth: Maximum depth to traverse (optional).
+            user_id: Optional user ID for trace operations.
+            tenant_id: Tenant ID (defaults to "1").
+            project_id: Optional project ID.
+            seller_id: Optional seller ID.
+            client_id: Optional client ID.
+            seller_profile_id: Optional seller profile ID.
+            raw_response: If True, return raw JSON response instead of parsed TraceNode.
+
+        Returns:
+            TraceNode representing the root of the sub-tree, or raw JSON dict if raw_response=True.
+        """
+        # Use trace_id if provided, otherwise use tree_id
+        final_tree_id = trace_id or tree_id
+        if not final_tree_id:
+            raise ValueError("Either trace_id or tree_id must be provided")
+        
+        return get_trace_tree(
+            api_key=self.api_key,
+            tree_id=final_tree_id,
+            node_id=node_id,
+            max_depth=max_depth,
+            user_id=user_id,
+            tenant_id=tenant_id,
+            project_id=project_id,
+            seller_id=seller_id,
+            client_id=client_id,
+            seller_profile_id=seller_profile_id,
+            raw_response=raw_response,
+        )
+
+    # ------------------------------------------------------------------
+    # SCHEDULING
+    # ------------------------------------------------------------------
+
+    async def schedule_run(
+        self,
+        workflow_id: str,
+        scheduled_run_type: str,
+        cron_expression: str,
+        time_zone: str,
+        input_mapping: List[Dict[str, str]],
+        file_url: Optional[str] = None,
+        batch_size: Optional[int] = None,
+    ) -> Dict[str, str]:
+        return await schedule_run_func(
+            workflow_id=workflow_id,
+            scheduled_run_type=scheduled_run_type,
+            cron_expression=cron_expression,
+            time_zone=time_zone,
+            input_mapping=input_mapping,
+            api_key=self.api_key,
+            file_url=file_url,
+            batch_size=batch_size,
+        )
+
+    async def cancel_schedule_run(
+        self,
+        job_master_id: Union[int, str],
+        unique_id: str,
+    ) -> Dict[str, str]:
+        """
+        Cancel a scheduled run by job_master_id and unique_id.
+        """
+        return await cancel_schedule_run_func(
+            job_master_id=job_master_id,
+            unique_id=unique_id,
+            api_key=self.api_key,
+        )
+
+    # ------------------------------------------------------------------
+    # BILLING
+    # ------------------------------------------------------------------
+
+    async def get_billing_detail(self, trace_id: str) -> Dict[str, Any]:
+        return await get_billing_detail(trace_id=trace_id, api_key=self.api_key)