niitaka-sdk 0.1.0__tar.gz

niitaka_sdk-0.1.0/.gitignore

@@ -0,0 +1,30 @@
+# Python
+__pycache__/
+*.pyc
+*.pyo
+.venv/
+venv/
+*.egg-info/
+
+# Environment / secrets — never commit these
+.env
+.env.local
+.env.*.local
+agents/.env
+backend/.env
+backend/.env.dev
+backend/.env.prod
+
+# Node / Next.js
+node_modules/
+.next/
+*.tsbuildinfo
+frontend/next-env.d.ts
+frontend/package-lock.json
+
+# macOS
+.DS_Store
+
+# IDE
+.idea/
+.vscode/

niitaka_sdk-0.1.0/PKG-INFO

@@ -0,0 +1,114 @@
+Metadata-Version: 2.4
+Name: niitaka-sdk
+Version: 0.1.0
+Summary: AI agent observability and control plane SDK
+Project-URL: Homepage, https://niitaka.ai
+Project-URL: Docs, https://docs.niitaka.ai
+Project-URL: Repository, https://github.com/amigotomodachiami/Niitaka
+Author-email: Niitaka <hello@niitaka.ai>
+License: MIT
+Keywords: agents,ai,anthropic,gemini,groq,langchain,llm,observability,openai,tracing
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: httpx>=0.24
+Requires-Dist: requests>=2.28
+Provides-Extra: all
+Requires-Dist: anthropic>=0.20; extra == 'all'
+Requires-Dist: google-generativeai>=0.4; extra == 'all'
+Requires-Dist: groq>=0.4; extra == 'all'
+Requires-Dist: langchain-core>=0.1; extra == 'all'
+Requires-Dist: openai>=1.0; extra == 'all'
+Provides-Extra: anthropic
+Requires-Dist: anthropic>=0.20; extra == 'anthropic'
+Provides-Extra: gemini
+Requires-Dist: google-generativeai>=0.4; extra == 'gemini'
+Provides-Extra: groq
+Requires-Dist: groq>=0.4; extra == 'groq'
+Provides-Extra: langchain
+Requires-Dist: langchain-core>=0.1; extra == 'langchain'
+Provides-Extra: openai
+Requires-Dist: openai>=1.0; extra == 'openai'
+Description-Content-Type: text/markdown
+
+# Niitaka SDK
+
+**AI agent observability and control plane — one import away.**
+
+Niitaka instruments your AI agents to capture every LLM call, tool invocation, and decision trace. Connect to the [Niitaka dashboard](https://niitaka.ai) to run experiments, enforce guardrails, and version your agents — without changing your agent code.
+
+## Install
+
+```bash
+pip install niitaka-sdk                    # core
+pip install "niitaka-sdk[openai]"          # + OpenAI
+pip install "niitaka-sdk[anthropic]"       # + Anthropic
+pip install "niitaka-sdk[all]"             # all providers
+```
+
+## Quick start
+
+```python
+import niitaka
+import os
+
+niitaka.configure(
+    api_key=os.getenv("NIITAKA_API_KEY"),
+    api_url="https://api.niitaka.ai",      # or your self-hosted URL
+)
+
+# Instrument your LLM provider (pick one or more)
+niitaka.instrument_openai()
+niitaka.instrument_anthropic()
+niitaka.instrument_gemini()
+niitaka.instrument_groq()
+
+# Wrap your agent run in a session
+with niitaka.start_session(goal="Summarise this document", agent_id="my-agent"):
+    response = client.chat.completions.create(...)  # auto-logged
+```
+
+Every call inside `start_session` is automatically traced — latency, token usage, cost, errors, and tool calls.
+
+## Providers
+
+| Provider | Extra | Import |
+|---|---|---|
+| OpenAI | `niitaka-sdk[openai]` | `niitaka.instrument_openai()` |
+| Anthropic | `niitaka-sdk[anthropic]` | `niitaka.instrument_anthropic()` |
+| Google Gemini | `niitaka-sdk[gemini]` | `niitaka.instrument_gemini()` |
+| Groq | `niitaka-sdk[groq]` | `niitaka.instrument_groq()` |
+| LangChain | `niitaka-sdk[langchain]` | `from niitaka import NiitakaCallbackHandler` |
+
+## LangChain
+
+```python
+from niitaka import NiitakaCallbackHandler
+import niitaka
+
+with niitaka.start_session(goal="...", agent_id="lc-agent"):
+    chain.invoke({"input": "..."}, config={"callbacks": [NiitakaCallbackHandler()]})
+```
+
+## Runtime config
+
+Niitaka can push model, temperature, and guardrail config to your agent at runtime — no redeploy needed:
+
+```python
+config = niitaka.get_runtime_config(agent_id="my-agent")
+# config["llm"]["model"], config["guardrails"]["cost_limit_usd"], ...
+```
+
+## Links
+
+- **Dashboard**: [niitaka.ai](https://niitaka.ai)
+- **Docs**: [docs.niitaka.ai](https://docs.niitaka.ai)
+- **Support**: hello@niitaka.ai

niitaka_sdk-0.1.0/README_SDK.md

@@ -0,0 +1,73 @@
+# Niitaka SDK
+
+**AI agent observability and control plane — one import away.**
+
+Niitaka instruments your AI agents to capture every LLM call, tool invocation, and decision trace. Connect to the [Niitaka dashboard](https://niitaka.ai) to run experiments, enforce guardrails, and version your agents — without changing your agent code.
+
+## Install
+
+```bash
+pip install niitaka-sdk                    # core
+pip install "niitaka-sdk[openai]"          # + OpenAI
+pip install "niitaka-sdk[anthropic]"       # + Anthropic
+pip install "niitaka-sdk[all]"             # all providers
+```
+
+## Quick start
+
+```python
+import niitaka
+import os
+
+niitaka.configure(
+    api_key=os.getenv("NIITAKA_API_KEY"),
+    api_url="https://api.niitaka.ai",      # or your self-hosted URL
+)
+
+# Instrument your LLM provider (pick one or more)
+niitaka.instrument_openai()
+niitaka.instrument_anthropic()
+niitaka.instrument_gemini()
+niitaka.instrument_groq()
+
+# Wrap your agent run in a session
+with niitaka.start_session(goal="Summarise this document", agent_id="my-agent"):
+    response = client.chat.completions.create(...)  # auto-logged
+```
+
+Every call inside `start_session` is automatically traced — latency, token usage, cost, errors, and tool calls.
+
+## Providers
+
+| Provider | Extra | Import |
+|---|---|---|
+| OpenAI | `niitaka-sdk[openai]` | `niitaka.instrument_openai()` |
+| Anthropic | `niitaka-sdk[anthropic]` | `niitaka.instrument_anthropic()` |
+| Google Gemini | `niitaka-sdk[gemini]` | `niitaka.instrument_gemini()` |
+| Groq | `niitaka-sdk[groq]` | `niitaka.instrument_groq()` |
+| LangChain | `niitaka-sdk[langchain]` | `from niitaka import NiitakaCallbackHandler` |
+
+## LangChain
+
+```python
+from niitaka import NiitakaCallbackHandler
+import niitaka
+
+with niitaka.start_session(goal="...", agent_id="lc-agent"):
+    chain.invoke({"input": "..."}, config={"callbacks": [NiitakaCallbackHandler()]})
+```
+
+## Runtime config
+
+Niitaka can push model, temperature, and guardrail config to your agent at runtime — no redeploy needed:
+
+```python
+config = niitaka.get_runtime_config(agent_id="my-agent")
+# config["llm"]["model"], config["guardrails"]["cost_limit_usd"], ...
+```
+
+## Links
+
+- **Dashboard**: [niitaka.ai](https://niitaka.ai)
+- **Docs**: [docs.niitaka.ai](https://docs.niitaka.ai)
+- **Support**: hello@niitaka.ai

niitaka_sdk-0.1.0/niitaka/__init__.py

@@ -0,0 +1,26 @@
+from .instrumentation import (
+    instrument_openai,
+    instrument_anthropic,
+    instrument_gemini,
+    instrument_groq,
+)
+from .client import get_runtime_config, start_session, configure
+
+# NiitakaCallbackHandler is optional — only available when langchain-core is installed.
+# Import it directly: from sdk.langchain_handler import NiitakaCallbackHandler
+try:
+    from .langchain_handler import NiitakaCallbackHandler
+except ImportError:
+    pass
+
+
+def auto_instrument():
+    """Instrument all installed LLM providers.
+
+    Calls each provider's instrument function and silently skips any that
+    are not installed.  Equivalent to calling each function individually.
+    """
+    instrument_openai()
+    instrument_anthropic()
+    instrument_gemini()
+    instrument_groq()

niitaka_sdk-0.1.0/niitaka/cli.py

@@ -0,0 +1,278 @@
+"""Niitaka CLI — manage policy files.
+
+Usage:
+    python -m sdk.cli push <policy-file> [--url URL] [--key KEY] [--dry-run]
+    python -m sdk.cli pull [--agent AGENT_ID] [--url URL] [--key KEY] [-o FILE]
+    python -m sdk.cli validate <policy-file>
+    python -m sdk.cli history [--agent AGENT_ID] [--limit N] [--url URL] [--key KEY]
+
+Commands
+--------
+push        Upload policies from a YAML/JSON file to the backend DB.
+            By default replaces all existing policies for each agent in the file.
+
+pull        Download current DB policies for one or all agents and print as YAML.
+            Use -o to write to a file instead of stdout.
+
+validate    Parse and validate a policy file without contacting the backend.
+
+history     Show the policy audit log — who changed what and when.
+            Use --agent to filter to a specific agent.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+import requests
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _headers(key: str) -> dict:
+    h = {"Content-Type": "application/json"}
+    if key:
+        h["X-API-Key"] = key
+    return h
+
+
+def _base_url(url: str) -> str:
+    return url.rstrip("/")
+
+
+# ── Commands ──────────────────────────────────────────────────────────────────
+
+def cmd_validate(args: argparse.Namespace) -> int:
+    from sdk.policy_loader import load_policy_file, PolicyFileError
+    try:
+        policies = load_policy_file(args.file)
+    except PolicyFileError as exc:
+        print(f"[ERROR] {exc}", file=sys.stderr)
+        return 1
+
+    from collections import Counter
+    counts = Counter(p["agent_id"] for p in policies)
+    print(f"✓ Valid policy file — {len(policies)} policies across {len(counts)} agent(s):")
+    for agent_id, n in sorted(counts.items()):
+        print(f"   {agent_id}: {n} rule(s)")
+    return 0
+
+
+def cmd_push(args: argparse.Namespace) -> int:
+    from sdk.policy_loader import load_policy_file, group_by_agent, PolicyFileError
+
+    # Validate file first
+    try:
+        policies = load_policy_file(args.file)
+    except PolicyFileError as exc:
+        print(f"[ERROR] {exc}", file=sys.stderr)
+        return 1
+
+    grouped = group_by_agent(policies)
+    base = _base_url(args.url)
+    headers = _headers(args.key)
+
+    print(f"[niitaka push] {args.file} → {base}")
+    print(f"  Agents: {', '.join(sorted(grouped))}")
+
+    if args.dry_run:
+        print("  [dry-run] No changes made.")
+        return 0
+
+    # Build import payload: list of policy objects with agent_id + config
+    payload_policies = []
+    for p in policies:
+        payload_policies.append({
+            "agent_id":  p["agent_id"],
+            "type":      p["type"],
+            "config": {
+                "condition": p["condition"],
+                "action":    p["action"],
+                "priority":  p["priority"],
+            },
+        })
+
+    try:
+        resp = requests.post(
+            f"{base}/policies/import",
+            headers=headers,
+            json={"policies": payload_policies, "replace": True},
+            timeout=15,
+        )
+        resp.raise_for_status()
+    except requests.HTTPError as exc:
+        print(f"[ERROR] Backend returned {exc.response.status_code}: {exc.response.text}",
+              file=sys.stderr)
+        return 1
+    except requests.RequestException as exc:
+        print(f"[ERROR] Could not reach backend at {base}: {exc}", file=sys.stderr)
+        return 1
+
+    result = resp.json()
+    print(f"  ✓ Imported {result.get('imported', '?')} policies "
+          f"(replaced {result.get('deleted', '?')} existing)")
+    for agent_id, n in sorted(result.get("by_agent", {}).items()):
+        print(f"     {agent_id}: {n} rule(s) written")
+    return 0
+
+
+def cmd_pull(args: argparse.Namespace) -> int:
+    from sdk.policy_loader import to_yaml_string
+
+    base = _base_url(args.url)
+    headers = _headers(args.key)
+
+    if args.agent:
+        # Single agent
+        try:
+            resp = requests.get(f"{base}/policies/{args.agent}", headers=headers, timeout=10)
+            resp.raise_for_status()
+        except requests.RequestException as exc:
+            print(f"[ERROR] {exc}", file=sys.stderr)
+            return 1
+        raw = resp.json()
+        # Attach agent_id (not in per-agent response)
+        for p in raw:
+            p["agent_id"] = args.agent
+    else:
+        # All policies
+        try:
+            resp = requests.get(f"{base}/policies", headers=headers, timeout=10)
+            resp.raise_for_status()
+        except requests.RequestException as exc:
+            print(f"[ERROR] {exc}", file=sys.stderr)
+            return 1
+        raw = resp.json()
+
+    if not raw:
+        print("# No policies found.", file=sys.stderr)
+        return 0
+
+    yaml_text = to_yaml_string(raw)
+
+    if args.output:
+        with open(args.output, "w", encoding="utf-8") as fh:
+            fh.write(yaml_text)
+        print(f"[niitaka pull] Wrote {len(raw)} policies to {args.output}")
+    else:
+        sys.stdout.write(yaml_text)
+
+    return 0
+
+
+def cmd_history(args: argparse.Namespace) -> int:
+    base    = _base_url(args.url)
+    headers = _headers(args.key)
+
+    url = f"{base}/policies/audit"
+    if args.agent:
+        url = f"{base}/policies/audit/{args.agent}"
+    url += f"?limit={args.limit}"
+
+    try:
+        resp = requests.get(url, headers=headers, timeout=10)
+        resp.raise_for_status()
+    except requests.RequestException as exc:
+        print(f"[ERROR] {exc}", file=sys.stderr)
+        return 1
+
+    entries = resp.json()
+    if not entries:
+        print("No audit log entries found.")
+        return 0
+
+    # Column widths
+    W_TIME   = 20
+    W_AGENT  = 20
+    W_ACTION = 8
+    W_TYPE   = 12
+    W_SOURCE = 6
+
+    header = (
+        f"{'Timestamp':<{W_TIME}}  "
+        f"{'Agent':<{W_AGENT}}  "
+        f"{'Action':<{W_ACTION}}  "
+        f"{'Type':<{W_TYPE}}  "
+        f"{'Src':<{W_SOURCE}}"
+    )
+    print(header)
+    print("─" * len(header))
+
+    for e in entries:
+        ts = e.get("created_at", "")
+        # Trim to readable format: "2026-04-12T14:32:01"
+        if ts and "T" in ts:
+            ts = ts[:19].replace("T", " ")
+
+        action  = e.get("action", "")
+        agent   = (e.get("agent_id") or "")[:W_AGENT]
+        ptype   = (e.get("policy_type") or "—")[:W_TYPE]
+        source  = (e.get("source") or "")[:W_SOURCE]
+
+        print(
+            f"{ts:<{W_TIME}}  "
+            f"{agent:<{W_AGENT}}  "
+            f"{action:<{W_ACTION}}  "
+            f"{ptype:<{W_TYPE}}  "
+            f"{source:<{W_SOURCE}}"
+        )
+
+    print(f"\n{len(entries)} entries")
+    return 0
+
+
+# ── Argument parser ───────────────────────────────────────────────────────────
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m sdk.cli",
+        description="Niitaka policy file manager",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    # Shared backend args
+    def add_backend_args(p: argparse.ArgumentParser) -> None:
+        p.add_argument("--url", default=os.getenv("NIITAKA_BACKEND_URL", "http://localhost:8000"),
+                       help="Backend URL (default: $NIITAKA_BACKEND_URL or http://localhost:8000)")
+        p.add_argument("--key", default=os.getenv("NIITAKA_API_KEY", ""),
+                       help="API key (default: $NIITAKA_API_KEY)")
+
+    # validate
+    p_val = sub.add_parser("validate", help="Validate a policy file without hitting the backend")
+    p_val.add_argument("file", help="Path to policy YAML or JSON file")
+
+    # push
+    p_push = sub.add_parser("push", help="Upload policy file to backend DB")
+    p_push.add_argument("file", help="Path to policy YAML or JSON file")
+    p_push.add_argument("--dry-run", action="store_true",
+                        help="Validate and show what would be pushed, but make no changes")
+    add_backend_args(p_push)
+
+    # pull
+    p_pull = sub.add_parser("pull", help="Download policies from backend as YAML")
+    p_pull.add_argument("--agent", default=None, help="Filter to a single agent_id")
+    p_pull.add_argument("-o", "--output", default=None, help="Write to file instead of stdout")
+    add_backend_args(p_pull)
+
+    # history
+    p_hist = sub.add_parser("history", help="Show the policy audit log")
+    p_hist.add_argument("--agent", default=None, help="Filter to a single agent_id")
+    p_hist.add_argument("--limit", type=int, default=50,
+                        help="Max entries to show (default: 50)")
+    add_backend_args(p_hist)
+
+    return parser
+
+
+def main() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+
+    dispatch = {"validate": cmd_validate, "push": cmd_push, "pull": cmd_pull, "history": cmd_history}
+    sys.exit(dispatch[args.command](args))
+
+
+if __name__ == "__main__":
+    main()