causaliq-knowledge 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

causaliq_knowledge/cli.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import json
 import sys
-from typing import Optional
+from typing import Any, Optional
 
 import click
 
@@ -405,6 +405,349 @@ def list_models(provider: Optional[str]) -> None:
     click.echo()
 
 
+# ============================================================================
+# Cache Commands
+# ============================================================================
+
+
+@cli.group("cache")
+def cache_group() -> None:
+    """Manage the LLM response cache.
+
+    Commands for inspecting, exporting, and importing cached LLM responses.
+
+    Examples:
+
+        cqknow cache stats ./llm_cache.db
+
+        cqknow cache export ./llm_cache.db ./export_dir
+
+        cqknow cache import ./llm_cache.db ./import_dir
+    """
+    pass
+
+
+@cache_group.command("stats")
+@click.argument("cache_path", type=click.Path(exists=True))
+@click.option(
+    "--json",
+    "output_json",
+    is_flag=True,
+    help="Output result as JSON.",
+)
+def cache_stats(cache_path: str, output_json: bool) -> None:
+    """Show cache statistics.
+
+    CACHE_PATH is the path to the SQLite cache database.
+
+    Examples:
+
+        cqknow cache stats ./llm_cache.db
+
+        cqknow cache stats ./llm_cache.db --json
+    """
+    from causaliq_knowledge.cache import TokenCache
+
+    try:
+        with TokenCache(cache_path) as cache:
+            entry_count = cache.entry_count()
+            token_count = cache.token_count()
+
+            if output_json:
+                output = {
+                    "cache_path": cache_path,
+                    "entry_count": entry_count,
+                    "token_count": token_count,
+                }
+                click.echo(json.dumps(output, indent=2))
+            else:
+                click.echo(f"\nCache: {cache_path}")
+                click.echo("=" * 40)
+                click.echo(f"Entries:  {entry_count:,}")
+                click.echo(f"Tokens:   {token_count:,}")
+                click.echo()
+    except Exception as e:
+        click.echo(f"Error opening cache: {e}", err=True)
+        sys.exit(1)
+
+
+@cache_group.command("export")
+@click.argument("cache_path", type=click.Path(exists=True))
+@click.argument("output_dir", type=click.Path())
+@click.option(
+    "--json",
+    "output_json",
+    is_flag=True,
+    help="Output result as JSON.",
+)
+def cache_export(cache_path: str, output_dir: str, output_json: bool) -> None:
+    """Export cache entries to human-readable files.
+
+    CACHE_PATH is the path to the SQLite cache database.
+    OUTPUT_DIR is the directory or zip file where files will be written.
+
+    If OUTPUT_DIR ends with .zip, entries are exported to a zip archive.
+    Otherwise, entries are exported to a directory.
+
+    Files are named using a human-readable format:
+        {model}_{node_a}_{node_b}_edge_{hash}.json
+
+    Examples:
+
+        cqknow cache export ./llm_cache.db ./export_dir
+
+        cqknow cache export ./llm_cache.db ./export.zip
+
+        cqknow cache export ./llm_cache.db ./export_dir --json
+    """
+    import tempfile
+    import zipfile
+    from pathlib import Path
+
+    from causaliq_knowledge.cache import TokenCache
+    from causaliq_knowledge.llm.cache import LLMCacheEntry, LLMEntryEncoder
+
+    output_path = Path(output_dir)
+    is_zip = output_path.suffix.lower() == ".zip"
+
+    try:
+        with TokenCache(cache_path) as cache:
+            # Register encoders for decoding
+            encoder = LLMEntryEncoder()
+            cache.register_encoder("llm", encoder)
+
+            # Register generic JsonEncoder for other types
+            from causaliq_knowledge.cache.encoders import JsonEncoder
+
+            json_encoder = JsonEncoder()
+            cache.register_encoder("json", json_encoder)
+
+            # Get entry types in the cache
+            entry_types = cache.list_entry_types()
+
+            if not entry_types:
+                if output_json:
+                    click.echo(json.dumps({"exported": 0, "error": None}))
+                else:
+                    click.echo("No entries to export.")
+                return
+
+            # Determine export directory (temp if zipping)
+            if is_zip:
+                temp_dir = tempfile.mkdtemp()
+                export_dir = Path(temp_dir)
+            else:
+                export_dir = output_path
+                export_dir.mkdir(parents=True, exist_ok=True)
+
+            # Export entries
+            exported = 0
+            for entry_type in entry_types:
+                if entry_type == "llm":
+                    # Query all entries of this type
+                    cursor = cache.conn.execute(
+                        "SELECT hash, data FROM cache_entries "
+                        "WHERE entry_type = ?",
+                        (entry_type,),
+                    )
+                    for cache_key, blob in cursor:
+                        data = encoder.decode(blob, cache)
+                        entry = LLMCacheEntry.from_dict(data)
+                        filename = encoder.generate_export_filename(
+                            entry, cache_key
+                        )
+                        file_path = export_dir / filename
+                        encoder.export_entry(entry, file_path)
+                        exported += 1
+                else:
+                    # For non-LLM types, use generic export
+                    count = cache.export_entries(export_dir, entry_type)
+                    exported += count
+
+            # Create zip archive if requested
+            if is_zip:
+                output_path.parent.mkdir(parents=True, exist_ok=True)
+                with zipfile.ZipFile(
+                    output_path, "w", zipfile.ZIP_DEFLATED
+                ) as zf:
+                    for file_path in export_dir.iterdir():
+                        if file_path.is_file():
+                            zf.write(file_path, file_path.name)
+                # Clean up temp directory
+                import shutil
+
+                shutil.rmtree(temp_dir)
+
+            # Output results
+            if output_json:
+                output = {
+                    "cache_path": cache_path,
+                    "output_path": str(output_path),
+                    "format": "zip" if is_zip else "directory",
+                    "exported": exported,
+                    "entry_types": entry_types,
+                }
+                click.echo(json.dumps(output, indent=2))
+            else:
+                fmt = "zip archive" if is_zip else "directory"
+                click.echo(
+                    f"\nExported {exported} entries to {fmt}: {output_path}"
+                )
+                click.echo(f"Entry types: {', '.join(entry_types)}")
+                click.echo()
+
+    except Exception as e:
+        click.echo(f"Error exporting cache: {e}", err=True)
+        sys.exit(1)
+
+
+def _is_llm_entry(data: Any) -> bool:
+    """Check if JSON data represents an LLM cache entry.
+
+    LLM entries have a specific structure with cache_key containing
+    model and messages, plus a response object.
+    """
+    if not isinstance(data, dict):
+        return False
+    cache_key = data.get("cache_key", {})
+    return (
+        isinstance(cache_key, dict)
+        and "model" in cache_key
+        and "messages" in cache_key
+        and "response" in data
+    )
+
+
+@cache_group.command("import")
+@click.argument("cache_path", type=click.Path())
+@click.argument("input_path", type=click.Path(exists=True))
+@click.option(
+    "--json",
+    "output_json",
+    is_flag=True,
+    help="Output result as JSON.",
+)
+def cache_import(cache_path: str, input_path: str, output_json: bool) -> None:
+    """Import cache entries from files.
+
+    CACHE_PATH is the path to the SQLite cache database (created if needed).
+    INPUT_PATH is a directory or zip file containing JSON files to import.
+
+    Entry types are auto-detected from JSON structure:
+    - LLM entries: contain cache_key.model, cache_key.messages, response
+    - Generic JSON: anything else
+
+    Examples:
+
+        cqknow cache import ./llm_cache.db ./import_dir
+
+        cqknow cache import ./llm_cache.db ./export.zip
+
+        cqknow cache import ./llm_cache.db ./import_dir --json
+    """
+    import hashlib
+    import tempfile
+    import zipfile
+    from pathlib import Path
+
+    from causaliq_knowledge.cache import TokenCache
+    from causaliq_knowledge.cache.encoders import JsonEncoder
+    from causaliq_knowledge.llm.cache import LLMEntryEncoder
+
+    input_file = Path(input_path)
+    is_zip = input_file.suffix.lower() == ".zip"
+
+    try:
+        with TokenCache(cache_path) as cache:
+            # Register encoders
+            llm_encoder = LLMEntryEncoder()
+            json_encoder = JsonEncoder()
+            cache.register_encoder("llm", llm_encoder)
+            cache.register_encoder("json", json_encoder)
+
+            # Determine input directory
+            if is_zip:
+                temp_dir = tempfile.mkdtemp()
+                import_dir = Path(temp_dir)
+                with zipfile.ZipFile(input_file, "r") as zf:
+                    zf.extractall(import_dir)
+            else:
+                import_dir = input_file
+                temp_dir = None
+
+            # Import entries
+            imported = 0
+            llm_count = 0
+            json_count = 0
+            skipped = 0
+
+            for file_path in import_dir.iterdir():
+                if (
+                    not file_path.is_file()
+                    or file_path.suffix.lower() != ".json"
+                ):
+                    continue
+
+                try:
+                    data = json.loads(file_path.read_text(encoding="utf-8"))
+                except (json.JSONDecodeError, UnicodeDecodeError):
+                    skipped += 1
+                    continue
+
+                # Detect entry type and generate cache key
+                if _is_llm_entry(data):
+                    # LLM entry - generate hash from cache_key contents
+                    cache_key_data = data.get("cache_key", {})
+                    key_str = json.dumps(cache_key_data, sort_keys=True)
+                    cache_key = hashlib.sha256(key_str.encode()).hexdigest()[
+                        :16
+                    ]
+                    cache.put_data(cache_key, "llm", data)
+                    llm_count += 1
+                else:
+                    # Generic JSON - use filename stem as key
+                    cache_key = file_path.stem
+                    cache.put_data(cache_key, "json", data)
+                    json_count += 1
+
+                imported += 1
+
+            # Clean up temp directory
+            if temp_dir:
+                import shutil
+
+                shutil.rmtree(temp_dir)
+
+            # Output results
+            if output_json:
+                output = {
+                    "cache_path": cache_path,
+                    "input_path": str(input_file),
+                    "format": "zip" if is_zip else "directory",
+                    "imported": imported,
+                    "llm_entries": llm_count,
+                    "json_entries": json_count,
+                    "skipped": skipped,
+                }
+                click.echo(json.dumps(output, indent=2))
+            else:
+                fmt = "zip archive" if is_zip else "directory"
+                click.echo(
+                    f"\nImported {imported} entries from {fmt}: {input_file}"
+                )
+                if llm_count:
+                    click.echo(f"  LLM entries: {llm_count}")
+                if json_count:
+                    click.echo(f"  JSON entries: {json_count}")
+                if skipped:
+                    click.echo(f"  Skipped: {skipped}")
+                click.echo()
+
+    except Exception as e:
+        click.echo(f"Error importing cache: {e}", err=True)
+        sys.exit(1)
+
+
 def main() -> None:
     """Entry point for the CLI."""
     cli()

causaliq_knowledge/llm/base_client.py

@@ -5,11 +5,18 @@ must implement. This provides a consistent API regardless of the
 underlying LLM provider.
 """
 
+from __future__ import annotations
+
+import hashlib
 import json
 import logging
+import time
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+if TYPE_CHECKING:  # pragma: no cover
+    from causaliq_knowledge.cache import TokenCache
 
 logger = logging.getLogger(__name__)
 
@@ -218,3 +225,136 @@ class BaseLLMClient(ABC):
             Model identifier string.
         """
         return getattr(self, "config", LLMConfig(model="unknown")).model
+
+    def _build_cache_key(
+        self,
+        messages: List[Dict[str, str]],
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+    ) -> str:
+        """Build a deterministic cache key for the request.
+
+        Creates a SHA-256 hash from the model, messages, temperature, and
+        max_tokens. The hash is truncated to 16 hex characters (64 bits).
+
+        Args:
+            messages: List of message dicts with "role" and "content" keys.
+            temperature: Sampling temperature (defaults to config value).
+            max_tokens: Maximum tokens (defaults to config value).
+
+        Returns:
+            16-character hex string cache key.
+        """
+        config = getattr(self, "config", LLMConfig(model="unknown"))
+        key_data = {
+            "model": config.model,
+            "messages": messages,
+            "temperature": (
+                temperature if temperature is not None else config.temperature
+            ),
+            "max_tokens": (
+                max_tokens if max_tokens is not None else config.max_tokens
+            ),
+        }
+        key_json = json.dumps(key_data, sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(key_json.encode()).hexdigest()[:16]
+
+    def set_cache(
+        self,
+        cache: Optional["TokenCache"],
+        use_cache: bool = True,
+    ) -> None:
+        """Configure caching for this client.
+
+        Args:
+            cache: TokenCache instance for caching, or None to disable.
+            use_cache: Whether to use the cache (default True).
+        """
+        self._cache = cache
+        self._use_cache = use_cache
+
+    @property
+    def cache(self) -> Optional["TokenCache"]:
+        """Return the configured cache, if any."""
+        return getattr(self, "_cache", None)
+
+    @property
+    def use_cache(self) -> bool:
+        """Return whether caching is enabled."""
+        return getattr(self, "_use_cache", True)
+
+    def cached_completion(
+        self,
+        messages: List[Dict[str, str]],
+        **kwargs: Any,
+    ) -> LLMResponse:
+        """Make a completion request with caching.
+
+        If caching is enabled and a cached response exists, returns
+        the cached response without making an API call. Otherwise,
+        makes the API call and caches the result.
+
+        Args:
+            messages: List of message dicts with "role" and "content" keys.
+            **kwargs: Provider-specific options (temperature, max_tokens, etc.)
+
+        Returns:
+            LLMResponse with the generated content and metadata.
+        """
+        from causaliq_knowledge.llm.cache import LLMCacheEntry, LLMEntryEncoder
+
+        cache = self.cache
+        use_cache = self.use_cache
+
+        # Build cache key
+        temperature = kwargs.get("temperature")
+        max_tokens = kwargs.get("max_tokens")
+        cache_key = self._build_cache_key(messages, temperature, max_tokens)
+
+        # Check cache
+        if use_cache and cache is not None:
+            # Ensure encoder is registered
+            if not cache.has_encoder("llm"):
+                cache.register_encoder("llm", LLMEntryEncoder())
+
+            if cache.exists(cache_key, "llm"):
+                cached_data = cache.get_data(cache_key, "llm")
+                if cached_data is not None:
+                    entry = LLMCacheEntry.from_dict(cached_data)
+                    return LLMResponse(
+                        content=entry.response.content,
+                        model=entry.model,
+                        input_tokens=entry.metadata.tokens.input,
+                        output_tokens=entry.metadata.tokens.output,
+                        cost=entry.metadata.cost_usd or 0.0,
+                    )
+
+        # Make API call with timing
+        start_time = time.perf_counter()
+        response = self.completion(messages, **kwargs)
+        latency_ms = int((time.perf_counter() - start_time) * 1000)
+
+        # Store in cache
+        if use_cache and cache is not None:
+            config = getattr(self, "config", LLMConfig(model="unknown"))
+            entry = LLMCacheEntry.create(
+                model=config.model,
+                messages=messages,
+                content=response.content,
+                temperature=(
+                    temperature
+                    if temperature is not None
+                    else config.temperature
+                ),
+                max_tokens=(
+                    max_tokens if max_tokens is not None else config.max_tokens
+                ),
+                provider=self.provider_name,
+                latency_ms=latency_ms,
+                input_tokens=response.input_tokens,
+                output_tokens=response.output_tokens,
+                cost_usd=response.cost,
+            )
+            cache.put_data(cache_key, "llm", entry.to_dict())
+
+        return response