basion-agent 0.4.0__py3-none-any.whl

basion_agent/structural/artifact.py

@@ -0,0 +1,94 @@
+"""Artifact structural streamer for file/image/iframe artifacts."""
+
+from typing import Optional, Dict, Any
+from .base import StructuralStreamer
+import uuid
+
+
+class Artifact(StructuralStreamer):
+    """Manages artifact lifecycle: generating → done.
+
+    Artifacts are files, images, iframes that are generated and displayed.
+    Artifact data is persisted to database on done().
+
+    Example:
+        artifact = Artifact()  # ID auto-generated
+        s.stream_by(artifact).generating("Creating chart...", progress=0.5)
+        s.stream_by(artifact).done(
+            url="https://example.com/chart.png",
+            type="image",
+            title="Sales Chart"
+        )
+    """
+
+    def _generate_id(self) -> str:
+        return f"artifact-{uuid.uuid4().hex[:8]}"
+
+    def get_event_type_prefix(self) -> str:
+        return "artifact"
+
+    def generating(
+        self,
+        message: Optional[str] = None,
+        progress: Optional[float] = None
+    ) -> 'Artifact':
+        """Signal artifact is being generated.
+
+        Args:
+            message: Status message (e.g., "Creating chart...")
+            progress: Progress value 0.0-1.0
+
+        Returns:
+            self for method chaining
+        """
+        kwargs = {}
+        if message is not None:
+            kwargs["message"] = message
+        if progress is not None:
+            kwargs["progress"] = progress
+
+        return self._send_event("generating", **kwargs)
+
+    def done(
+        self,
+        url: str,
+        type: str,
+        title: Optional[str] = None,
+        description: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> 'Artifact':
+        """Signal artifact is complete and persist to database.
+
+        Args:
+            url: URL to artifact
+            type: Artifact type (image, iframe, document, video, audio, code, link, file)
+            title: Display title
+            description: Optional description
+            metadata: Additional metadata (size, mimeType, dimensions, etc.)
+
+        Returns:
+            self for method chaining
+        """
+        kwargs = {
+            "url": url,
+            "type": type,
+        }
+        if title:
+            kwargs["title"] = title
+        if description:
+            kwargs["description"] = description
+        if metadata:
+            kwargs["metadata"] = metadata
+
+        return self._send_event("done", **kwargs)
+
+    def error(self, message: str) -> 'Artifact':
+        """Signal artifact generation failed.
+
+        Args:
+            message: Error message
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("error", message=message)

basion_agent/structural/base.py

@@ -0,0 +1,71 @@
+"""Base class for structural streaming objects."""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, TYPE_CHECKING
+import uuid
+
+if TYPE_CHECKING:
+    from ..streamer import Streamer
+
+
+class StructuralStreamer(ABC):
+    """Base class for structural streaming objects.
+
+    Structural streamers generate events but don't hold streaming state.
+    The Streamer binds to them via stream_by() to enable streaming.
+    """
+
+    def __init__(self, id: Optional[str] = None):
+        """Initialize with auto-generated or explicit ID."""
+        self.id = id or self._generate_id()
+        self._streamer: Optional['Streamer'] = None
+
+    @abstractmethod
+    def _generate_id(self) -> str:
+        """Generate a unique ID for this streamer type."""
+        pass
+
+    @abstractmethod
+    def get_event_type_prefix(self) -> str:
+        """Get the event type prefix (e.g., 'artifact', 'text_block')."""
+        pass
+
+    def _bind_streamer(self, streamer: 'Streamer') -> 'StructuralStreamer':
+        """Bind this structural streamer to a Streamer instance.
+
+        Called by Streamer.stream_by(). Returns self for chaining.
+        """
+        self._streamer = streamer
+        return self
+
+    def _send_event(self, action: str, **kwargs) -> 'StructuralStreamer':
+        """Send an event through the bound streamer.
+
+        Creates event payload and sends via streamer._send().
+        Returns self for method chaining.
+        """
+        if not self._streamer:
+            raise RuntimeError(
+                f"{self.__class__.__name__} not bound to a Streamer. "
+                f"Call streamer.stream_by() first."
+            )
+
+        # Build event body
+        body = {"id": self.id, **kwargs}
+
+        # Build event type
+        event_type = f"{self.get_event_type_prefix()}:{action}"
+
+        # Send through streamer using existing _send() method
+        # content = stringified JSON body
+        # eventType = structural event type
+        # persist = False (structural events not persisted as content)
+        import json
+        self._streamer._send(
+            content=json.dumps(body),
+            done=False,
+            persist=False,
+            event_type=event_type
+        )
+
+        return self

basion_agent/structural/stepper.py

@@ -0,0 +1,125 @@
+"""Stepper structural streamer for multi-step progress indicators."""
+
+from typing import Optional, List
+from .base import StructuralStreamer
+import uuid
+
+
+class Stepper(StructuralStreamer):
+    """Manages multi-step progress indicator.
+
+    Stepper events are NOT persisted to database.
+    IDs are auto-generated - no need to provide them manually.
+
+    Example:
+        stepper = Stepper(steps=["Fetch", "Process", "Report"])  # ID auto-generated
+        s.stream_by(stepper).start_step(0)
+        s.stream_by(stepper).complete_step(0)
+        s.stream_by(stepper).start_step(1)
+        s.stream_by(stepper).add_step("Verify")  # Dynamic step
+        s.stream_by(stepper).complete_step(1)
+        s.stream_by(stepper).done()
+    """
+
+    def __init__(
+        self,
+        id: Optional[str] = None,
+        steps: Optional[List[str]] = None
+    ):
+        super().__init__(id)
+        self.steps = steps or []
+        self._initialized = False
+
+    def _generate_id(self) -> str:
+        return f"stepper-{uuid.uuid4().hex[:8]}"
+
+    def get_event_type_prefix(self) -> str:
+        return "stepper"
+
+    def _ensure_initialized(self) -> 'Stepper':
+        """Send initialization event with steps if not already sent.
+
+        This is automatically called before any other stepper event.
+        """
+        if not self._initialized and self.steps:
+            self._send_event("init", steps=self.steps)
+            self._initialized = True
+        return self
+
+    def start_step(self, index: int) -> 'Stepper':
+        """Mark step as in progress.
+
+        Args:
+            index: Step index to start
+
+        Returns:
+            self for method chaining
+        """
+        self._ensure_initialized()
+        return self._send_event("start_step", index=index)
+
+    def complete_step(self, index: int) -> 'Stepper':
+        """Mark step as complete.
+
+        Args:
+            index: Step index to complete
+
+        Returns:
+            self for method chaining
+        """
+        self._ensure_initialized()
+        return self._send_event("complete_step", index=index)
+
+    def fail_step(self, index: int, error: Optional[str] = None) -> 'Stepper':
+        """Mark step as failed.
+
+        Args:
+            index: Step index that failed
+            error: Optional error message
+
+        Returns:
+            self for method chaining
+        """
+        self._ensure_initialized()
+        kwargs = {"index": index}
+        if error:
+            kwargs["error"] = error
+        return self._send_event("fail_step", **kwargs)
+
+    def add_step(self, label: str, index: Optional[int] = None) -> 'Stepper':
+        """Add a new step dynamically.
+
+        Args:
+            label: Step label
+            index: Optional position to insert at
+
+        Returns:
+            self for method chaining
+        """
+        self._ensure_initialized()
+        kwargs = {"label": label}
+        if index is not None:
+            kwargs["index"] = index
+        return self._send_event("add_step", **kwargs)
+
+    def update_step_label(self, index: int, label: str) -> 'Stepper':
+        """Update step label.
+
+        Args:
+            index: Step index to update
+            label: New label
+
+        Returns:
+            self for method chaining
+        """
+        self._ensure_initialized()
+        return self._send_event("update_step_label", index=index, label=label)
+
+    def done(self) -> 'Stepper':
+        """Mark stepper as complete (all steps finished).
+
+        Returns:
+            self for method chaining
+        """
+        self._ensure_initialized()
+        return self._send_event("done")

basion_agent/structural/surface.py

@@ -0,0 +1,90 @@
+"""Surface structural streamer for embedded interactive components."""
+
+from typing import Optional, Dict, Any
+from .base import StructuralStreamer
+import uuid
+
+
+class Surface(StructuralStreamer):
+    """Manages embedded surface lifecycle (identical to Artifact).
+
+    Surfaces are embedded interactive components (iframes, embeds).
+    Surface data is persisted to database as artifacts with surface type.
+
+    Example:
+        surface = Surface()
+        s.stream_by(surface).generating("Loading calendar...")
+        s.stream_by(surface).done(url="https://cal.com/embed", type="iframe")
+    """
+
+    def _generate_id(self) -> str:
+        return f"surface-{uuid.uuid4().hex[:8]}"
+
+    def get_event_type_prefix(self) -> str:
+        return "surface"
+
+    def generating(
+        self,
+        message: Optional[str] = None,
+        progress: Optional[float] = None
+    ) -> 'Surface':
+        """Signal surface is loading.
+
+        Args:
+            message: Status message (e.g., "Loading calendar...")
+            progress: Progress value 0.0-1.0
+
+        Returns:
+            self for method chaining
+        """
+        kwargs = {}
+        if message is not None:
+            kwargs["message"] = message
+        if progress is not None:
+            kwargs["progress"] = progress
+
+        return self._send_event("generating", **kwargs)
+
+    def done(
+        self,
+        url: str,
+        type: str = "iframe",
+        title: Optional[str] = None,
+        description: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> 'Surface':
+        """Signal surface is ready and persist to database.
+
+        Args:
+            url: URL to surface
+            type: Surface type (default: iframe)
+            title: Display title
+            description: Optional description
+            metadata: Additional metadata
+
+        Returns:
+            self for method chaining
+        """
+        kwargs = {
+            "url": url,
+            "type": type,
+        }
+        if title:
+            kwargs["title"] = title
+        if description:
+            kwargs["description"] = description
+        if metadata:
+            kwargs["metadata"] = metadata
+
+        return self._send_event("done", **kwargs)
+
+    def error(self, message: str) -> 'Surface':
+        """Signal surface loading failed.
+
+        Args:
+            message: Error message
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("error", message=message)

basion_agent/structural/text_block.py

@@ -0,0 +1,96 @@
+"""TextBlock structural streamer for collapsible text content."""
+
+from typing import Optional
+from .base import StructuralStreamer
+import uuid
+
+
+class TextBlock(StructuralStreamer):
+    """Manages collapsible text block with title and body.
+
+    Both title and body support streaming (append) and update (replace).
+    TextBlock events are NOT persisted to database.
+
+    Example:
+        block = TextBlock()  # ID auto-generated
+        s.stream_by(block).set_variant("thinking")
+        s.stream_by(block).stream_title("Deep ")
+        s.stream_by(block).stream_title("Analysis...")
+        s.stream_by(block).stream_body("Step 1: Check patterns\\n")
+        s.stream_by(block).stream_body("Step 2: Validate\\n")
+        s.stream_by(block).update_title("Analysis Complete")
+        s.stream_by(block).done()
+    """
+
+    def __init__(self, id: Optional[str] = None, title: Optional[str] = None):
+        super().__init__(id)
+        self.initial_title = title
+
+    def _generate_id(self) -> str:
+        return f"text_block-{uuid.uuid4().hex[:8]}"
+
+    def get_event_type_prefix(self) -> str:
+        return "text_block"
+
+    def stream_title(self, content: str) -> 'TextBlock':
+        """Stream title content (appends to existing).
+
+        Args:
+            content: Content to append to title
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("stream_title", content=content)
+
+    def stream_body(self, content: str) -> 'TextBlock':
+        """Stream body content (appends to existing).
+
+        Args:
+            content: Content to append to body
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("stream_body", content=content)
+
+    def update_title(self, title: str) -> 'TextBlock':
+        """Replace entire title.
+
+        Args:
+            title: New title (replaces existing)
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("update_title", title=title)
+
+    def update_body(self, body: str) -> 'TextBlock':
+        """Replace entire body.
+
+        Args:
+            body: New body (replaces existing)
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("update_body", body=body)
+
+    def set_variant(self, variant: str) -> 'TextBlock':
+        """Set visual variant.
+
+        Args:
+            variant: Visual style (thinking, note, warning, error, success)
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("set_variant", variant=variant)
+
+    def done(self) -> 'TextBlock':
+        """Mark text block as complete.
+
+        Returns:
+            self for method chaining
+        """
+        return self._send_event("done")

basion_agent/tools/__init__.py

@@ -0,0 +1,19 @@
+"""Tools module for Basion Agent SDK."""
+
+from .container import Tools
+from .knowledge_graph import (
+    Edge,
+    Entity,
+    KnowledgeGraphTool,
+    PathStep,
+    SimilarDisease,
+)
+
+__all__ = [
+    "Tools",
+    "KnowledgeGraphTool",
+    "Entity",
+    "Edge",
+    "SimilarDisease",
+    "PathStep",
+]

basion_agent/tools/container.py

@@ -0,0 +1,46 @@
+"""Tools container - accessed via agent.tools property."""
+
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from .knowledge_graph import KnowledgeGraphTool
+
+
+class Tools:
+    """
+    Container for all available tools.
+
+    Accessed via agent.tools property (similar to agent.streamer(message)):
+        kg = agent.tools.knowledge_graph
+        diseases = await kg.search_diseases(name="Huntington")
+
+    Tools are lazy-initialized - no overhead if not used.
+    """
+
+    def __init__(self, gateway_client):
+        self._gateway_client = gateway_client
+        self._knowledge_graph: Optional["KnowledgeGraphTool"] = None
+
+    @property
+    def knowledge_graph(self) -> "KnowledgeGraphTool":
+        """
+        Knowledge Graph tool for querying biomedical entities and relationships.
+
+        Example:
+            kg = agent.tools.knowledge_graph
+            diseases = await kg.search_diseases(name="Huntington")
+            similar = await kg.find_similar_diseases("Huntington Disease")
+            path = await kg.find_shortest_path("BRCA1", "protein", "Breast Cancer", "disease")
+        """
+        if self._knowledge_graph is None:
+            from .knowledge_graph import KnowledgeGraphTool
+
+            self._knowledge_graph = KnowledgeGraphTool(
+                self._gateway_client.knowledge_graph_url
+            )
+        return self._knowledge_graph
+
+    async def close(self):
+        """Close all tool sessions."""
+        if self._knowledge_graph:
+            await self._knowledge_graph.close()