augagent 0.1.0__py3-none-any.whl

augagent/models.py

@@ -0,0 +1,291 @@
+"""Core Pydantic models for the augagent framework.
+
+All data flowing through the system is strongly typed via Pydantic v2 models,
+providing runtime validation, serialisation, and rich editor support.
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, SecretStr
+
+
+# ---------------------------------------------------------------------------
+# Enumerations
+# ---------------------------------------------------------------------------
+
+class Role(str, Enum):
+    """Roles within a conversation turn."""
+
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+    TOOL = "tool"
+
+
+class TaskStatus(str, Enum):
+    """Lifecycle status of a :class:`Task`."""
+
+    PENDING = "pending"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# LLM CONFIGURATION SCHEMAS
+# ═══════════════════════════════════════════════════════════════════════════
+
+class LLMConfig(BaseModel):
+    """Declarative configuration for connecting to an LLM provider.
+
+    Supports any **OpenAI-compatible** API (OpenAI, Azure OpenAI, Ollama,
+    vLLM, LiteLLM, etc.) by pointing ``base_url`` at the right endpoint.
+
+    Example::
+
+        # Explicit key
+        config = LLMConfig(model="gpt-4o", api_key="sk-...")
+
+        # From environment (default: reads OPENAI_API_KEY)
+        config = LLMConfig(model="gpt-4o")
+
+        # Local Ollama
+        config = LLMConfig(
+            model="llama3",
+            base_url="http://localhost:11434/v1",
+            api_key="ollama",
+        )
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    model: str = Field(
+        default="gpt-4o",
+        description="Model identifier (e.g. 'gpt-4o', 'claude-3-opus', 'llama3').",
+    )
+    api_key: SecretStr | None = Field(
+        default=None,
+        description=(
+            "API key for the provider. If not set, the framework reads the "
+            "environment variable specified by ``api_key_env_var``."
+        ),
+    )
+    api_key_env_var: str = Field(
+        default="OPENAI_API_KEY",
+        description="Environment variable to read the API key from when ``api_key`` is None.",
+    )
+    base_url: str = Field(
+        default="https://api.openai.com/v1",
+        description="Base URL of the chat-completions API.",
+    )
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    max_tokens: int = Field(default=4096, ge=1)
+    top_p: float = Field(default=1.0, ge=0.0, le=1.0)
+    frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
+    presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
+    stop: list[str] | None = Field(
+        default=None,
+        description="Up to 4 sequences where the model will stop generating.",
+    )
+    timeout: float = Field(
+        default=120.0,
+        gt=0,
+        description="HTTP request timeout in seconds.",
+    )
+    max_retries: int = Field(
+        default=3,
+        ge=0,
+        description="Retries on transient HTTP errors (429, 500, 502, 503).",
+    )
+    extra_headers: dict[str, str] = Field(
+        default_factory=dict,
+        description="Additional HTTP headers sent with every request.",
+    )
+
+    def resolve_api_key(self) -> str:
+        """Return the API key, resolving from the environment if necessary.
+
+        Raises
+        ------
+        ValueError
+            If no key is available from either the field or the env var.
+        """
+        if self.api_key is not None:
+            return self.api_key.get_secret_value()
+        key = os.environ.get(self.api_key_env_var, "")
+        if not key:
+            raise ValueError(
+                f"No API key provided and environment variable "
+                f"'{self.api_key_env_var}' is not set."
+            )
+        return key
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# CHAT COMPLETION RESPONSE SCHEMAS  (OpenAI-compatible)
+# ═══════════════════════════════════════════════════════════════════════════
+
+class FunctionCall(BaseModel):
+    """The function the model wants to invoke (name + raw JSON args)."""
+
+    name: str
+    arguments: str  # JSON-encoded string — parsed by the caller
+
+
+class ChatToolCall(BaseModel):
+    """A single tool-call entry inside an assistant message."""
+
+    id: str
+    type: str = "function"
+    function: FunctionCall
+
+
+class ChatMessage(BaseModel):
+    """A message object as returned inside a ``Choice``."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    role: str
+    content: str | None = None
+    tool_calls: list[ChatToolCall] | None = None
+    refusal: str | None = None
+
+
+class Choice(BaseModel):
+    """One candidate completion returned by the API."""
+
+    index: int = 0
+    message: ChatMessage
+    finish_reason: str | None = None
+
+
+class TokenUsage(BaseModel):
+    """Token-level usage statistics from the API response."""
+
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+
+class ChatCompletion(BaseModel):
+    """Fully-parsed ``/chat/completions`` response.
+
+    Using a Pydantic model here means malformed API responses are caught
+    immediately with clear validation errors rather than producing silent
+    ``KeyError`` s downstream.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    id: str = ""
+    object: str = "chat.completion"
+    created: int | None = None
+    model: str = ""
+    choices: list[Choice] = Field(default_factory=list)
+    usage: TokenUsage | None = None
+
+
+# ---------------------------------------------------------------------------
+# Internal message models
+# ---------------------------------------------------------------------------
+
+class ToolCall(BaseModel):
+    """Framework-internal representation of a tool invocation request."""
+
+    model_config = ConfigDict(frozen=True)
+
+    id: str = Field(default_factory=lambda: f"call_{uuid.uuid4().hex[:12]}")
+    tool_name: str = Field(..., description="Name of the tool to invoke.")
+    arguments: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Keyword arguments for the tool.",
+    )
+
+
+class ToolResponse(BaseModel):
+    """Encapsulates the result returned by a tool invocation."""
+
+    model_config = ConfigDict(frozen=True)
+
+    tool_call_id: str = Field(..., description="ID of the originating ToolCall.")
+    tool_name: str
+    content: str = Field(..., description="Serialised output from the tool.")
+    is_error: bool = False
+
+
+class Message(BaseModel):
+    """A single message within an agent conversation.
+
+    Messages form the backbone of inter-agent and agent-LLM communication.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    id: str = Field(default_factory=lambda: uuid.uuid4().hex)
+    role: Role
+    content: str | None = None
+    tool_calls: list[ToolCall] = Field(default_factory=list)
+    tool_responses: list[ToolResponse] = Field(default_factory=list)
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+# ---------------------------------------------------------------------------
+# Agent configuration  (serialisable subset of AugAgent fields)
+# ---------------------------------------------------------------------------
+
+class AgentConfig(BaseModel):
+    """Declarative, serialisable configuration for an
+    :class:`~augagent.agent.AugAgent`.
+
+    Useful for persisting, templating, or sharing agent definitions
+    as JSON / YAML without carrying callable state.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    name: str = Field(..., min_length=1, max_length=128)
+    role: str = Field(..., description="A short descriptor of the agent's persona.")
+    goal: str = Field(..., description="What this agent is trying to achieve.")
+    backstory: str = Field(
+        default="",
+        description="Rich context that shapes the agent's behaviour.",
+    )
+    llm_config: LLMConfig = Field(default_factory=LLMConfig)
+    max_iterations: int = Field(
+        default=25,
+        ge=1,
+        description="Maximum ReAct loop iterations before the agent stops.",
+    )
+    allow_delegation: bool = Field(
+        default=False,
+        description="Whether this agent may delegate sub-tasks to teammates.",
+    )
+    verbose: bool = False
+
+
+# ---------------------------------------------------------------------------
+# Task result
+# ---------------------------------------------------------------------------
+
+class TaskResult(BaseModel):
+    """The output produced when a :class:`~augagent.task.Task` completes."""
+
+    model_config = ConfigDict(frozen=True)
+
+    task_id: str
+    agent_name: str
+    status: TaskStatus = TaskStatus.COMPLETED
+    output: str = ""
+    raw_output: Any = None
+    token_usage: dict[str, int] = Field(default_factory=dict)
+    elapsed_seconds: float = 0.0
+    iterations: int = 0
+    metadata: dict[str, Any] = Field(default_factory=dict)

augagent/task.py

@@ -0,0 +1,156 @@
+"""Task — a discrete unit of work assigned to an AugAgent.
+
+Tasks describe *what* needs to be done, *what* the expected output looks
+like, and *which* agent is responsible.  They are the currency of work
+within an :class:`~augagent.team.AugTeam`.
+
+Example::
+
+    from augagent import AugAgent, AugTask, LLMConfig
+
+    writer = AugAgent(
+        name="Writer",
+        role="Content Writer",
+        goal="Write great copy",
+        llm_config=LLMConfig(model="gpt-4o"),
+    )
+
+    task = AugTask(
+        description="Write a blog post about AI agents",
+        expected_output="A 500-word blog post in Markdown format",
+        agent=writer,
+    )
+"""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from augagent.agent import AugAgent
+from augagent.models import TaskResult, TaskStatus
+from augagent.telemetry import get_logger
+
+
+class AugTask(BaseModel):
+    """A unit of work to be executed by an :class:`AugAgent`.
+
+    Parameters
+    ----------
+    description:
+        A natural-language description of what needs to be done.
+    expected_output:
+        A description of the ideal output format and content.
+    agent:
+        The agent responsible for executing this task.
+    context:
+        Optional list of predecessor :class:`AugTask` instances whose outputs
+        are prepended to this task's prompt, enabling information flow
+        across sequential tasks.
+    async_execution:
+        If ``True``, the task is eligible for concurrent execution.
+    output_json:
+        Optional Pydantic model class to validate and parse the output.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: str = Field(default_factory=lambda: uuid.uuid4().hex[:12])
+    description: str = Field(..., min_length=1)
+    expected_output: str = Field(default="")
+    agent: AugAgent | None = None
+    context: list["AugTask"] = Field(default_factory=list)
+    async_execution: bool = False
+    output_json: Any = Field(default=None, exclude=True, repr=False)
+
+    # -- runtime state (excluded from serialisation) ------------------------
+    status: TaskStatus = Field(default=TaskStatus.PENDING, exclude=True)
+    result: TaskResult | None = Field(default=None, exclude=True)
+
+    # -- public API ---------------------------------------------------------
+
+    async def execute(self, agent: AugAgent | None = None) -> TaskResult:
+        """Run the task using its assigned agent (or an override).
+
+        Raises
+        ------
+        ValueError
+            If no agent is assigned or provided.
+        """
+        executor = agent or self.agent
+        if executor is None:
+            raise ValueError(
+                f"Task '{self.id}' has no assigned agent. "
+                "Pass an agent explicitly or set task.agent."
+            )
+
+        logger = get_logger()
+        self.status = TaskStatus.IN_PROGRESS
+
+        # Build the full prompt, incorporating context from upstream tasks
+        prompt = self._build_prompt()
+
+        try:
+            result = await executor.execute(prompt)
+            # Re-stamp with this task's id
+            result = result.model_copy(update={"task_id": self.id})
+            self.result = result
+            self.status = TaskStatus.COMPLETED
+
+            # Optional JSON output validation
+            if self.output_json is not None:
+                try:
+                    self.output_json.model_validate_json(result.output)
+                except Exception as exc:
+                    logger.log_error(f"JSON validation failed: {exc}")
+
+            return result
+
+        except Exception as exc:
+            self.status = TaskStatus.FAILED
+            logger.log_error(f"Task failed: {exc}")
+            return TaskResult(
+                task_id=self.id,
+                agent_name=executor.name,
+                status=TaskStatus.FAILED,
+                output=f"Task failed: {exc}",
+            )
+
+    # -- internals ----------------------------------------------------------
+
+    def _build_prompt(self) -> str:
+        """Compose the task prompt, prepending upstream context if present."""
+        parts: list[str] = []
+
+        if self.context:
+            parts.append("## Context from previous tasks\n")
+            for ctx_task in self.context:
+                if ctx_task.result and ctx_task.result.output:
+                    parts.append(
+                        f"### Output from '{ctx_task.description}'\n{ctx_task.result.output}\n"
+                    )
+            parts.append("---\n")
+
+        parts.append(f"## Task\n{self.description}\n")
+        if self.expected_output:
+            parts.append(f"## Expected Output\n{self.expected_output}\n")
+
+        return "\n".join(parts)
+
+    def interpolate_inputs(self, inputs: dict[str, Any]) -> None:
+        """Interpolate inputs into description and expected_output."""
+        self.description = self.description.format(**inputs)
+        if self.expected_output:
+            self.expected_output = self.expected_output.format(**inputs)
+
+    def __repr__(self) -> str:
+        agent_name = self.agent.name if self.agent else "unassigned"
+        return (
+            f"AugTask(id={self.id!r}, agent={agent_name!r}, "
+            f"status={self.status.value!r})"
+        )
+
+# Backward compatibility alias
+Task = AugTask

augagent/team.py

@@ -0,0 +1,163 @@
+"""Team — multi-agent orchestration.
+
+An :class:`AugTeam` groups agents and tasks together. In sequential mode,
+it executes tasks one by one, automatically chaining the output of one task
+as context for the next.
+
+Example::
+
+    from augagent import AugAgent, AugTask, AugTeam, LLMConfig
+
+    config = LLMConfig(model="gpt-4o")
+    researcher = AugAgent(name="Researcher", role="Researcher", goal="Find facts", llm_config=config)
+    writer     = AugAgent(name="Writer", role="Writer", goal="Write articles", llm_config=config)
+
+    t1 = AugTask(description="Research AI trends", expected_output="Bullet points", agent=researcher)
+    t2 = AugTask(description="Write blog post", expected_output="Markdown article", agent=writer)
+
+    # Output from t1 is automatically passed as context to t2 in sequential mode
+    team = AugTeam(
+        agents=[researcher, writer],
+        tasks=[t1, t2],
+    )
+    results = team.kickoff()
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from augagent.agent import AugAgent
+from augagent.models import TaskResult, TaskStatus
+from augagent.task import AugTask
+from augagent.telemetry import get_logger
+
+
+class Process(str, Enum):
+    """Orchestration strategy for a :class:`AugTeam`."""
+
+    SEQUENTIAL = "sequential"
+
+
+class AugTeam(BaseModel):
+    """A coordinated group of :class:`AugAgent` instances executing :class:`AugTask` objects.
+
+    Parameters
+    ----------
+    agents:
+        The roster of agents available to this team.
+    tasks:
+        An ordered list of tasks to execute.
+    process:
+        The orchestration strategy (currently only ``"sequential"``).
+    verbose:
+        Enable detailed logging.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    agents: list[AugAgent] = Field(..., min_length=1)
+    tasks: list[AugTask] = Field(..., min_length=1)
+    process: Process = Process.SEQUENTIAL
+    verbose: bool = False
+
+    # -- public API ---------------------------------------------------------
+
+    def kickoff(self, inputs: dict[str, Any] | None = None) -> list[TaskResult]:
+        """Execute all tasks sequentially and return results.
+
+        Passes the output of each completed task as context to the next task.
+        Handles both sync and async contexts transparently.
+        """
+        if inputs:
+            for task in self.tasks:
+                task.interpolate_inputs(inputs)
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+
+        if loop and loop.is_running():
+            import concurrent.futures
+
+            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+                return pool.submit(asyncio.run, self._run()).result()
+        return asyncio.run(self._run())
+
+    async def akickoff(self, inputs: dict[str, Any] | None = None) -> list[TaskResult]:
+        """Async version of :meth:`kickoff`."""
+        if inputs:
+            for task in self.tasks:
+                task.interpolate_inputs(inputs)
+        return await self._run()
+
+    # -- orchestration strategies -------------------------------------------
+
+    async def _run(self) -> list[TaskResult]:
+        logger = get_logger()
+        start = time.time()
+
+        if self.verbose:
+            logger.log_info(f"Team kickoff started with {len(self.tasks)} tasks.")
+
+        results = await self._run_sequential(logger)
+
+        elapsed = time.time() - start
+        completed = sum(1 for r in results if r.status == TaskStatus.COMPLETED)
+        failed = sum(1 for r in results if r.status == TaskStatus.FAILED)
+
+        if self.verbose:
+            logger.log_info(f"Team finished in {elapsed:.2f}s ({completed} completed, {failed} failed).")
+
+        return results
+
+    async def _run_sequential(self, logger: Any) -> list[TaskResult]:
+        """Execute tasks one-by-one, passing outputs as context to the next."""
+        results: list[TaskResult] = []
+
+        for i, task in enumerate(self.tasks):
+            # Automatically chain context: if this isn't the first task, and the previous task
+            # successfully generated output, append it to this task's context.
+            if i > 0 and results[-1].status == TaskStatus.COMPLETED:
+                prev_task = self.tasks[i - 1]
+                if prev_task not in task.context:
+                    task.context.append(prev_task)
+
+            executor = task.agent or self.agents[0]
+
+            logger.log_handoff(from_agent="Manager", to_agent=executor.name, task_desc=task.description)
+
+            result = await task.execute(agent=executor)
+            results.append(result)
+
+            if result.status == TaskStatus.COMPLETED:
+                logger.log_info(f"Task completed successfully. Output length: {len(result.output)} chars.")
+            else:
+                logger.log_error(f"Task failed: {result.output}")
+
+        return results
+
+    # -- utilities ----------------------------------------------------------
+
+    def get_agent(self, name: str) -> AugAgent | None:
+        """Look up an agent by name."""
+        for agent in self.agents:
+            if agent.name == name:
+                return agent
+        return None
+
+    def __repr__(self) -> str:
+        return (
+            f"AugTeam(agents={[a.name for a in self.agents]}, "
+            f"tasks={len(self.tasks)}, process={self.process.value!r})"
+        )
+
+
+# Backward compatibility alias
+Team = AugTeam

augagent/telemetry.py

@@ -0,0 +1,50 @@
+"""Telemetry & logging primitives.
+
+Provides a lightweight logger that prints agent handoffs and tool executions
+to the console in a clean, readable format.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+
+# Configure a basic console logger
+_logger = logging.getLogger("augagent")
+_logger.setLevel(logging.INFO)
+
+if not _logger.handlers:
+    handler = logging.StreamHandler(sys.stdout)
+    formatter = logging.Formatter("%(message)s")
+    handler.setFormatter(formatter)
+    _logger.addHandler(handler)
+
+
+class AgentLogger:
+    """A lightweight logger for tracking agent orchestration."""
+
+    def __init__(self, logger: logging.Logger = _logger):
+        self._logger = logger
+
+    def log_handoff(self, from_agent: str, to_agent: str, task_desc: str) -> None:
+        """Log a task handoff to an agent."""
+        self._logger.info(f"\n[{from_agent} -> {to_agent}] Starting Task: {task_desc}")
+
+    def log_tool_execution(self, agent_name: str, tool_name: str, args: dict[str, Any] | str) -> None:
+        """Log when an agent executes a tool."""
+        self._logger.info(f"  [Tool Execution] {agent_name} is using '{tool_name}' with args: {args}")
+
+    def log_info(self, message: str) -> None:
+        """Log a general info message."""
+        self._logger.info(f"  [Info] {message}")
+
+    def log_error(self, message: str) -> None:
+        """Log an error message."""
+        self._logger.error(f"  [Error] {message}")
+
+
+_DEFAULT_LOGGER = AgentLogger()
+
+def get_logger() -> AgentLogger:
+    """Return the global agent logger instance."""
+    return _DEFAULT_LOGGER