citekit 0.1.0__tar.gz

citekit-0.1.0/PKG-INFO

@@ -0,0 +1,45 @@
+Metadata-Version: 2.4
+Name: citekit
+Version: 0.1.0
+Summary: A local SDK that lets AI agents open specific parts of files instead of sending entire documents.
+Project-URL: Repository, https://github.com/citekit/citekit
+Author: CiteKit Team
+License: MIT
+Keywords: ai,audio,llm,mcp,pdf,rag,sdk,video
+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.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: click>=8.0
+Requires-Dist: google-genai>=1.0
+Requires-Dist: mcp>=1.0
+Requires-Dist: pillow>=10.0
+Requires-Dist: pydantic>=2.0
+Requires-Dist: pymupdf>=1.24
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# CiteKit SDK for Python
+
+See the main [README](../../README.md) for usage instructions.
+
+## Installation
+
+```bash
+pip install citekit
+```
+
+## Quick Start
+
+```python
+from citekit import CiteKitClient, GeminiMapper
+
+# ... (see main docs)
+```

citekit-0.1.0/README.md

@@ -0,0 +1,17 @@
+# CiteKit SDK for Python
+
+See the main [README](../../README.md) for usage instructions.
+
+## Installation
+
+```bash
+pip install citekit
+```
+
+## Quick Start
+
+```python
+from citekit import CiteKitClient, GeminiMapper
+
+# ... (see main docs)
+```

citekit-0.1.0/citekit/__init__.py

@@ -0,0 +1,21 @@
+"""CiteKit — Let AI agents open specific parts of files."""
+
+from citekit.models import ResourceMap, Node, Location, ResolvedEvidence
+from citekit.client import CiteKitClient
+from citekit.address import parse_address, build_address
+from citekit.mapper.base import MapperProvider
+from citekit.mapper.gemini import GeminiMapper
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "ResourceMap",
+    "Node",
+    "Location",
+    "ResolvedEvidence",
+    "CiteKitClient",
+    "parse_address",
+    "build_address",
+    "MapperProvider",
+    "GeminiMapper",
+]

citekit-0.1.0/citekit/__main__.py

@@ -0,0 +1,4 @@
+from citekit.cli import main
+
+if __name__ == "__main__":
+    main()

citekit-0.1.0/citekit/address.py

@@ -0,0 +1,166 @@
+"""Address parser and builder for CiteKit URI-style addresses.
+
+Formats:
+    doc://resource_id#pages=3-5
+    video://resource_id#t=192-230
+    audio://resource_id#t=60-120
+    image://resource_id#bbox=0.2,0.3,0.8,0.7
+"""
+
+from __future__ import annotations
+
+import re
+from urllib.parse import urlparse, parse_qs
+
+from citekit.models import Location
+
+# Scheme → modality mapping
+_SCHEME_TO_MODALITY = {
+    "doc": "document",
+    "video": "video",
+    "audio": "audio",
+    "image": "image",
+}
+
+_MODALITY_TO_SCHEME = {v: k for k, v in _SCHEME_TO_MODALITY.items()}
+
+
+def parse_address(address: str) -> tuple[str, Location]:
+    """Parse a CiteKit address into (resource_id, Location).
+
+    Examples:
+        >>> parse_address("doc://calculus_book#pages=12-13")
+        ('calculus_book', Location(modality='document', pages=[12, 13]))
+
+        >>> parse_address("video://lecture1#t=192-230")
+        ('lecture1', Location(modality='video', start=192.0, end=230.0))
+
+        >>> parse_address("image://diagram#bbox=0.2,0.3,0.8,0.7")
+        ('diagram', Location(modality='image', bbox=(0.2, 0.3, 0.8, 0.7)))
+    """
+    # Parse scheme manually since urllib doesn't handle custom schemes well
+    match = re.match(r"^(\w+)://([^#]+)(?:#(.+))?$", address)
+    if not match:
+        raise ValueError(f"Invalid CiteKit address: {address}")
+
+    scheme, resource_id, fragment = match.group(1), match.group(2), match.group(3)
+
+    if scheme not in _SCHEME_TO_MODALITY:
+        raise ValueError(f"Unknown scheme '{scheme}'. Expected one of: {list(_SCHEME_TO_MODALITY.keys())}")
+
+    modality = _SCHEME_TO_MODALITY[scheme]
+
+    # Parse fragment parameters
+    pages = None
+    start = None
+    end = None
+    bbox = None
+
+    if fragment:
+        params = dict(part.split("=", 1) for part in fragment.split("&") if "=" in part)
+
+        if "pages" in params:
+            page_str = params["pages"]
+            if "-" in page_str:
+                p_start, p_end = page_str.split("-", 1)
+                pages = list(range(int(p_start), int(p_end) + 1))
+            else:
+                pages = [int(p) for p in page_str.split(",")]
+
+        if "t" in params:
+            time_str = params["t"]
+            if "-" in time_str:
+                t_start, t_end = time_str.split("-", 1)
+                start = _parse_time(t_start)
+                end = _parse_time(t_end)
+
+        if "bbox" in params:
+            parts = params["bbox"].split(",")
+            if len(parts) != 4:
+                raise ValueError(f"bbox must have 4 values, got {len(parts)}")
+            bbox = tuple(float(p) for p in parts)
+
+    return resource_id, Location(
+        modality=modality,
+        pages=pages,
+        start=start,
+        end=end,
+        bbox=bbox,
+    )
+
+
+def build_address(resource_id: str, location: Location) -> str:
+    """Build a CiteKit address from a resource ID and location.
+
+    Examples:
+        >>> build_address("book", Location(modality="document", pages=[3, 4, 5]))
+        'doc://book#pages=3-5'
+
+        >>> build_address("lecture", Location(modality="video", start=192.0, end=230.0))
+        'video://lecture#t=192-230'
+    """
+    scheme = _MODALITY_TO_SCHEME.get(location.modality)
+    if not scheme:
+        raise ValueError(f"Unknown modality: {location.modality}")
+
+    fragment_parts = []
+
+    if location.pages is not None:
+        if len(location.pages) == 0:
+            raise ValueError("Pages list cannot be empty")
+        pages_sorted = sorted(location.pages)
+        # Check if pages are consecutive for range notation
+        if pages_sorted == list(range(pages_sorted[0], pages_sorted[-1] + 1)):
+            fragment_parts.append(f"pages={pages_sorted[0]}-{pages_sorted[-1]}")
+        else:
+            fragment_parts.append(f"pages={','.join(str(p) for p in pages_sorted)}")
+
+    if location.start is not None and location.end is not None:
+        start_str = _format_time(location.start)
+        end_str = _format_time(location.end)
+        fragment_parts.append(f"t={start_str}-{end_str}")
+
+    if location.bbox is not None:
+        bbox_str = ",".join(f"{v:g}" for v in location.bbox)
+        fragment_parts.append(f"bbox={bbox_str}")
+
+    fragment = "&".join(fragment_parts)
+    if fragment:
+        return f"{scheme}://{resource_id}#{fragment}"
+    return f"{scheme}://{resource_id}"
+
+
+def _parse_time(time_str: str) -> float:
+    """Parse a time string to seconds. Supports seconds or HH:MM:SS format."""
+    time_str = time_str.strip()
+
+    # HH:MM:SS or MM:SS format
+    if ":" in time_str:
+        parts = time_str.split(":")
+        if len(parts) == 3:
+            h, m, s = parts
+            return int(h) * 3600 + int(m) * 60 + float(s)
+        elif len(parts) == 2:
+            m, s = parts
+            return int(m) * 60 + float(s)
+
+    return float(time_str)
+
+
+def _format_time(seconds: float) -> str:
+    """Format seconds as a compact time string."""
+    if seconds == int(seconds):
+        seconds = int(seconds)
+
+    # Use HH:MM:SS for large values
+    if isinstance(seconds, int) and seconds >= 3600:
+        h = seconds // 3600
+        m = (seconds % 3600) // 60
+        s = seconds % 60
+        return f"{h:02d}:{m:02d}:{s:02d}"
+    elif isinstance(seconds, int) and seconds >= 60:
+        m = seconds // 60
+        s = seconds % 60
+        return f"{m:02d}:{s:02d}"
+
+    return str(seconds)

citekit-0.1.0/citekit/cli.py

@@ -0,0 +1,119 @@
+import asyncio
+import functools
+import json
+import os
+import sys
+from pathlib import Path
+
+import click
+from citekit.client import CiteKitClient
+
+def async_command(f):
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        return asyncio.run(f(*args, **kwargs))
+    return wrapper
+
+@click.group()
+def main():
+    """CiteKit CLI - Local AI Resource Mapper & Resolver."""
+    pass
+
+@main.command()
+@click.argument("path", type=click.Path(exists=True))
+@click.option("--type", "-t", help="Resource type (document, video, audio, image). If omitted, inferred from extension.")
+@async_command
+async def ingest(path, type):
+    """Ingest a file and generate a resource map."""
+    client = CiteKitClient()
+    
+    if not type:
+        ext = Path(path).suffix.lower()
+        if ext in (".pdf", ".txt", ".md"):
+            type = "document"
+        elif ext in (".mp4", ".mov", ".avi", ".mkv"):
+            type = "video"
+        elif ext in (".mp3", ".wav", ".m4a"):
+            type = "audio"
+        elif ext in (".png", ".jpg", ".jpeg", ".webp"):
+            type = "image"
+        else:
+            click.echo(f"⚠️ Could not infer type from extension '{ext}'. Please specify --type.", err=True)
+            sys.exit(1)
+
+    click.echo(f"🔍 Ingesting {path} as '{type}'...")
+    try:
+        resource_map = await client.ingest(path, resource_type=type)
+        click.echo(f"✅ Map generated: {resource_map.resource_id}")
+        click.echo(f"   Title: {resource_map.title}")
+        click.echo(f"   Nodes: {len(resource_map.nodes)}")
+    except Exception as e:
+        click.echo(f"❌ Error: {e}", err=True)
+        sys.exit(1)
+
+@main.command()
+@click.argument("node_id")
+@async_command
+async def resolve(node_id):
+    """Resolve a node ID to its value (file chunk)."""
+    client = CiteKitClient()
+    
+    click.echo(f"📎 Resolving node: {node_id}")
+    try:
+        evidence = await client.resolve(node_id)
+        click.echo(f"✅ Output: {evidence.output_path}")
+        click.echo(f"   Modality: {evidence.modality}")
+        click.echo(f"   Address: {evidence.address}")
+    except Exception as e:
+        click.echo(f"❌ Error: {e}", err=True)
+        sys.exit(1)
+
+@main.command()
+@click.argument("resource_id")
+@async_command
+async def structure(resource_id):
+    """Get the JSON structure (map) for a resource ID."""
+    client = CiteKitClient()
+    
+    try:
+        resource_map = await client.get_map(resource_id)
+        click.echo(json.dumps(resource_map.model_dump(), indent=2, default=str))
+    except Exception as e:
+        click.echo(f"❌ Error: {e}", err=True)
+        sys.exit(1)
+
+@main.command("list")
+@async_command
+async def list_resources():
+    """List all ingested resources."""
+    map_dir = Path(".resource_maps")
+    if not map_dir.exists():
+        click.echo("No resources found (directory .resource_maps missing).")
+        return
+
+    maps = list(map_dir.glob("*.json"))
+    if not maps:
+        click.echo("No resources found.")
+        return
+
+    click.echo(f"found {len(maps)} resources:")
+    for map_file in maps:
+        try:
+            data = json.loads(map_file.read_text(encoding="utf-8"))
+            click.echo(f" - {data.get('resource_id')} ({data.get('type')}): {data.get('title')}")
+        except:
+            click.echo(f" - {map_file.name} (corrupt)")
+
+@main.command()
+@async_command
+async def serve():
+    """Run the MCP server (stdio mode) for AI agents."""
+    from citekit.mcp_server import create_server
+    from mcp.server.stdio import stdio_server
+
+    server = create_server()
+    async with stdio_server() as (read_stream, write_stream):
+        await server.run(read_stream, write_stream, server.create_initialization_options())
+
+if __name__ == "__main__":
+    main()

citekit-0.1.0/citekit/client.py

@@ -0,0 +1,176 @@
+"""CiteKit client — main entry point for the SDK."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from citekit.address import build_address
+from citekit.mapper.base import MapperProvider
+from citekit.models import Location, Node, ResolvedEvidence, ResourceMap
+from citekit.resolvers.audio import AudioResolver
+from citekit.resolvers.document import DocumentResolver
+from citekit.resolvers.image import ImageResolver
+from citekit.resolvers.video import VideoResolver
+
+
+class CiteKitClient:
+    """Main client for ingesting resources, reading maps, and resolving nodes.
+
+    Usage:
+        from citekit import CiteKitClient, GeminiMapper
+
+        mapper = GeminiMapper(api_key="YOUR_KEY")
+        client = CiteKitClient(mapper=mapper)
+
+        # Ingest a PDF
+        resource_map = await client.ingest("textbook.pdf", "document")
+
+        # Later: resolve a specific node
+        evidence = client.resolve("textbook", "derivatives.definition")
+        print(evidence.output_path)  # → .citekit_output/textbook_pages_12-13.pdf
+    """
+
+    def __init__(
+        self,
+        mapper: MapperProvider | None = None,
+        storage_dir: str = ".resource_maps",
+        output_dir: str = ".citekit_output",
+    ):
+        self._mapper = mapper
+        self._storage_dir = Path(storage_dir)
+        self._output_dir = Path(output_dir)
+        self._storage_dir.mkdir(parents=True, exist_ok=True)
+        self._output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Initialize resolvers
+        self._resolvers = {
+            "document": DocumentResolver(output_dir=output_dir),
+            "video": VideoResolver(output_dir=output_dir),
+            "audio": AudioResolver(output_dir=output_dir),
+            "image": ImageResolver(output_dir=output_dir),
+        }
+
+    # ── Ingestion ────────────────────────────────────────────────────────────
+
+    async def ingest(
+        self,
+        resource_path: str,
+        resource_type: str,
+        resource_id: str | None = None,
+    ) -> ResourceMap:
+        """Ingest a resource: analyze it with the mapper and save the map locally.
+
+        Args:
+            resource_path: Path to the resource file (PDF, video, audio, image).
+            resource_type: One of "document", "video", "audio", "image".
+            resource_id: Optional custom ID. Defaults to the filename stem.
+
+        Returns:
+            The generated ResourceMap.
+        """
+        if self._mapper is None:
+            raise RuntimeError(
+                "No mapper provider configured. "
+                "Pass a MapperProvider (e.g. GeminiMapper) to CiteKitClient()."
+            )
+
+        resource_map = await self._mapper.generate_map(
+            resource_path=resource_path,
+            resource_type=resource_type,
+            resource_id=resource_id,
+        )
+
+        # Save to local storage
+        self._save_map(resource_map)
+
+        return resource_map
+
+    # ── Map access ───────────────────────────────────────────────────────────
+
+    def get_map(self, resource_id: str) -> ResourceMap:
+        """Load a previously generated map from local storage.
+
+        Args:
+            resource_id: The resource ID to look up.
+
+        Returns:
+            The ResourceMap.
+
+        Raises:
+            FileNotFoundError: If no map exists for this resource_id.
+        """
+        map_path = self._storage_dir / f"{resource_id}.json"
+        if not map_path.exists():
+            raise FileNotFoundError(
+                f"No map found for resource '{resource_id}'. "
+                f"Expected at: {map_path}"
+            )
+
+        data = json.loads(map_path.read_text(encoding="utf-8"))
+        return ResourceMap.model_validate(data)
+
+    def list_maps(self) -> list[str]:
+        """List all available resource map IDs."""
+        return [
+            p.stem
+            for p in self._storage_dir.glob("*.json")
+        ]
+
+    def get_structure(self, resource_id: str) -> dict:
+        """Get the map as a plain dict — useful for MCP/JSON responses."""
+        return self.get_map(resource_id).model_dump(mode="json")
+
+    # ── Resolution ───────────────────────────────────────────────────────────
+
+    def resolve(self, resource_id: str, node_id: str) -> ResolvedEvidence:
+        """Resolve a node into extracted evidence.
+
+        Args:
+            resource_id: The resource to look up.
+            node_id: The node ID within that resource.
+
+        Returns:
+            ResolvedEvidence with the path to the extracted file.
+        """
+        resource_map = self.get_map(resource_id)
+        node = resource_map.get_node(node_id)
+
+        if node is None:
+            available = resource_map.list_node_ids()
+            raise ValueError(
+                f"Node '{node_id}' not found in resource '{resource_id}'. "
+                f"Available nodes: {available}"
+            )
+
+        # Pick the right resolver
+        modality = node.location.modality
+        resolver = self._resolvers.get(modality)
+
+        if resolver is None:
+            raise ValueError(f"No resolver for modality: {modality}")
+
+        # Resolve
+        output_path = resolver.resolve(node, resource_map.source_path)
+
+        # Build the URI address
+        address = build_address(resource_id, node.location)
+
+        return ResolvedEvidence(
+            output_path=output_path,
+            modality=modality,
+            address=address,
+            node=node,
+            resource_id=resource_id,
+        )
+
+    # ── Private ──────────────────────────────────────────────────────────────
+
+    def _save_map(self, resource_map: ResourceMap) -> None:
+        """Save a ResourceMap to local JSON storage."""
+        map_path = self._storage_dir / f"{resource_map.resource_id}.json"
+        data = resource_map.model_dump(mode="json")
+        map_path.write_text(
+            json.dumps(data, indent=2, ensure_ascii=False),
+            encoding="utf-8",
+        )

citekit-0.1.0/citekit/mapper/__init__.py

@@ -0,0 +1 @@
+"""Mapper providers for CiteKit."""

citekit-0.1.0/citekit/mapper/base.py

@@ -0,0 +1,36 @@
+"""Abstract base class for mapper providers."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from citekit.models import ResourceMap
+
+
+class MapperProvider(ABC):
+    """Base class for all mapper providers.
+
+    A mapper analyzes a resource file and produces a structured ResourceMap
+    containing nodes that point to physical locations within the resource.
+
+    To create a custom provider, subclass this and implement generate_map().
+    """
+
+    @abstractmethod
+    async def generate_map(
+        self,
+        resource_path: str,
+        resource_type: str,
+        resource_id: str | None = None,
+    ) -> ResourceMap:
+        """Analyze a resource and produce a structured map.
+
+        Args:
+            resource_path: Path to the resource file.
+            resource_type: Type of resource ("document", "video", "audio", "image").
+            resource_id: Optional custom ID. If not provided, derived from filename.
+
+        Returns:
+            A ResourceMap with nodes pointing to locations in the resource.
+        """
+        ...