contexttrace 0.1.0__py3-none-any.whl

contexttrace/config.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+from contexttrace.errors import ContextTraceConfigError
+
+
+DEFAULT_PROJECT = "default"
+DEFAULT_BASE_URL = "http://localhost:8000"
+DEFAULT_MODE = "local"
+DEFAULT_LOCAL_STORE_DIR = ".contexttrace"
+DEFAULT_STORAGE_PATH = ".contexttrace/contexttrace.db"
+CONFIG_FILE = "contexttrace.yaml"
+
+
+@dataclass(frozen=True)
+class ContextTraceConfig:
+    api_key: Optional[str] = None
+    project: str = DEFAULT_PROJECT
+    base_url: str = DEFAULT_BASE_URL
+    mode: str = DEFAULT_MODE
+    local_only: bool = True
+    timeout: float = 30.0
+    retries: int = 2
+    debug: bool = False
+    local_store_dir: str = DEFAULT_LOCAL_STORE_DIR
+    storage_path: str = DEFAULT_STORAGE_PATH
+    log_chunk_text: bool = True
+    log_answer_text: bool = True
+    eval_endpoint: Optional[str] = None
+    judge_provider: str = "local"
+
+
+def load_config(
+    *,
+    api_key: Optional[str] = None,
+    project: Optional[str] = None,
+    base_url: Optional[str] = None,
+    api_url: Optional[str] = None,
+    mode: Optional[str] = None,
+    local_only: Optional[bool] = None,
+    timeout: Optional[float] = None,
+    retries: Optional[int] = None,
+    debug: Optional[bool] = None,
+    local_store_dir: Optional[str] = None,
+    storage_path: Optional[str] = None,
+    log_chunk_text: Optional[bool] = None,
+    log_answer_text: Optional[bool] = None,
+    eval_endpoint: Optional[str] = None,
+    judge_provider: Optional[str] = None,
+    config_path: Optional[str] = None,
+) -> ContextTraceConfig:
+    file_values = _read_config_file(config_path)
+    env_base_url = os.getenv("CONTEXTTRACE_API_URL") or os.getenv("CONTEXTTRACE_BASE_URL")
+    explicit_api_url = api_url or base_url
+    resolved_mode = _first(
+        mode,
+        os.getenv("CONTEXTTRACE_MODE"),
+        file_values.get("mode"),
+        "hosted" if explicit_api_url and local_only is not True else None,
+        "hosted" if env_base_url and local_only is not True else None,
+        DEFAULT_MODE,
+    )
+    resolved_local_store_dir = str(
+        _first(
+            local_store_dir,
+            os.getenv("CONTEXTTRACE_LOCAL_STORE_DIR"),
+            file_values.get("local_store_dir"),
+            DEFAULT_LOCAL_STORE_DIR,
+        )
+    )
+    resolved_storage_path = str(
+        _first(
+            storage_path,
+            os.getenv("CONTEXTTRACE_STORAGE_PATH"),
+            file_values.get("storage_path"),
+            str(Path(resolved_local_store_dir) / "contexttrace.db"),
+        )
+    )
+    resolved_local_only = _as_bool(
+        _first(
+            local_only,
+            os.getenv("CONTEXTTRACE_LOCAL_ONLY"),
+            file_values.get("local_only"),
+            False if resolved_mode == "hosted" else True,
+        )
+    )
+
+    resolved = ContextTraceConfig(
+        api_key=_first(
+            api_key,
+            os.getenv("CONTEXTTRACE_API_KEY"),
+            file_values.get("api_key"),
+        ),
+        project=str(
+            _first(
+                project,
+                os.getenv("CONTEXTTRACE_PROJECT"),
+                file_values.get("project"),
+                DEFAULT_PROJECT,
+            )
+        ),
+        base_url=str(
+            _first(
+                api_url,
+                base_url,
+                env_base_url,
+                file_values.get("base_url"),
+                file_values.get("api_url"),
+                DEFAULT_BASE_URL,
+            )
+        ),
+        mode=str(resolved_mode),
+        local_only=resolved_local_only,
+        timeout=float(
+            _first(
+                timeout,
+                os.getenv("CONTEXTTRACE_TIMEOUT"),
+                file_values.get("timeout"),
+                30.0,
+            )
+        ),
+        retries=int(
+            _first(
+                retries,
+                os.getenv("CONTEXTTRACE_RETRIES"),
+                file_values.get("retries"),
+                2,
+            )
+        ),
+        debug=_as_bool(
+            _first(
+                debug,
+                os.getenv("CONTEXTTRACE_DEBUG"),
+                file_values.get("debug"),
+                False,
+            )
+        ),
+        local_store_dir=resolved_local_store_dir,
+        storage_path=resolved_storage_path,
+        log_chunk_text=_as_bool(
+            _first(
+                log_chunk_text,
+                os.getenv("CONTEXTTRACE_LOG_CHUNK_TEXT"),
+                file_values.get("log_chunk_text"),
+                True,
+            )
+        ),
+        log_answer_text=_as_bool(
+            _first(
+                log_answer_text,
+                os.getenv("CONTEXTTRACE_LOG_ANSWER_TEXT"),
+                file_values.get("log_answer_text"),
+                True,
+            )
+        ),
+        eval_endpoint=_first(
+            eval_endpoint,
+            os.getenv("CONTEXTTRACE_EVAL_ENDPOINT"),
+            file_values.get("eval_endpoint"),
+        ),
+        judge_provider=str(
+            _first(
+                judge_provider,
+                os.getenv("CONTEXTTRACE_JUDGE_PROVIDER"),
+                file_values.get("judge_provider"),
+                "local",
+            )
+        ),
+    )
+
+    if resolved.mode not in {"hosted", "local"}:
+        raise ContextTraceConfigError("ContextTrace mode must be 'hosted' or 'local'.")
+    return resolved
+
+
+def write_default_config(path: str = CONFIG_FILE, *, overwrite: bool = False) -> str:
+    output = Path(path)
+    if output.exists() and not overwrite:
+        return str(output)
+    output.write_text(
+        "\n".join(
+            [
+                "mode: local",
+                "project: default",
+                "local_only: true",
+                "local_store_dir: .contexttrace",
+                "storage_path: .contexttrace/contexttrace.db",
+                "log_chunk_text: true",
+                "log_answer_text: true",
+                "judge_provider: local",
+                "api_key: ''",
+                "base_url: ''",
+                "timeout: 30",
+                "retries: 2",
+                "debug: false",
+                "eval_endpoint: ''",
+                "",
+            ]
+        ),
+        encoding="utf-8",
+    )
+    return str(output)
+
+
+def _read_config_file(config_path: Optional[str]) -> dict[str, Any]:
+    candidates = [Path(config_path)] if config_path else [Path(CONFIG_FILE)]
+    for candidate in candidates:
+        if candidate.exists():
+            return _parse_simple_yaml(candidate.read_text(encoding="utf-8"))
+    return {}
+
+
+def _parse_simple_yaml(text: str) -> dict[str, Any]:
+    values: dict[str, Any] = {}
+    for raw_line in text.splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#") or ":" not in line:
+            continue
+        key, value = line.split(":", 1)
+        parsed = value.strip().strip("\"'")
+        if parsed == "":
+            values[key.strip()] = None
+        elif parsed.lower() in {"true", "false"}:
+            values[key.strip()] = parsed.lower() == "true"
+        else:
+            values[key.strip()] = parsed
+    return values
+
+
+def _first(*values: Any) -> Any:
+    for value in values:
+        if value is not None:
+            return value
+    return None
+
+
+def _as_bool(value: Any) -> bool:
+    if isinstance(value, bool):
+        return value
+    if value is None:
+        return False
+    return str(value).strip().lower() in {"1", "true", "yes", "on"}

contexttrace/demo.py

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable
+
+from contexttrace.client import ContextTrace
+from contexttrace.demo_data import load_demo_dataset
+from contexttrace.report import ReportGenerator
+
+
+STRATEGY_TOP_K = {
+    "dense_top_k": 2,
+    "bm25": 2,
+    "bm25_top_k": 2,
+    "hybrid": 3,
+    "hybrid_rerank": 4,
+    "corrective": 4,
+    "corrective_rag": 4,
+    "adaptive": 4,
+    "contexttrace_adaptive": 4,
+}
+
+
+@dataclass(frozen=True)
+class DemoRun:
+    dataset: str
+    eval_run_id: str | None
+    trace_ids: list[str]
+    report_path: str
+    summary: dict[str, Any]
+
+
+def run_demo_dataset(
+    *,
+    dataset: str,
+    contexttrace: ContextTrace,
+    strategy: str = "adaptive",
+    report_path: str | None = None,
+    max_questions: int | None = None,
+) -> DemoRun:
+    loaded = load_demo_dataset(dataset)
+    chunks = _document_chunks(loaded["documents"])
+    questions = list(loaded["questions"])
+    if max_questions is not None:
+        questions = questions[:max_questions]
+
+    trace_ids: list[str] = []
+    traces: list[dict[str, Any]] = []
+    for index, question in enumerate(questions):
+        trace_id = _run_question(
+            contexttrace=contexttrace,
+            dataset_name=loaded["name"],
+            question=question,
+            expected_answer=loaded["expected_answers"].get(question["id"], ""),
+            expected_sources=list(loaded["expected_sources"].get(question["id"], [])),
+            chunks=chunks,
+            strategy=strategy,
+            position=index,
+        )
+        trace_ids.append(trace_id)
+        traces.append(contexttrace.get_trace(trace_id))
+
+    summary = aggregate_trace_metrics(traces)
+    store = getattr(getattr(contexttrace, "_transport", None), "store", None)
+    eval_run_id = None
+    if store is not None:
+        eval_run_id = store.create_eval_run(dataset=loaded["name"], endpoint="contexttrace-demo", summary=summary)
+        for index, question in enumerate(questions):
+            store.add_eval_question(
+                eval_run_id=eval_run_id,
+                question=question,
+                trace_id=trace_ids[index],
+                position=index,
+            )
+
+    if report_path is None:
+        report_path = str(Path(".contexttrace") / "reports" / ("%s_demo.html" % loaded["name"]))
+    ReportGenerator().generate_eval_report(
+        {
+            "id": eval_run_id or "%s-demo" % loaded["name"],
+            "dataset": loaded["name"],
+            "endpoint": "contexttrace-demo",
+            "summary": summary,
+        },
+        traces,
+        path=report_path,
+    )
+    return DemoRun(
+        dataset=loaded["name"],
+        eval_run_id=eval_run_id,
+        trace_ids=trace_ids,
+        report_path=report_path,
+        summary=summary,
+    )
+
+
+def aggregate_trace_metrics(traces: Iterable[dict[str, Any]]) -> dict[str, Any]:
+    rows = list(traces)
+    count = len(rows)
+    if not rows:
+        return {
+            "questions_tested": 0,
+            "reliability_score": 0.0,
+            "failure_rate": 0.0,
+            "citation_support": 0.0,
+            "unsupported_claim_rate": 0.0,
+            "retrieval_miss_rate": 0.0,
+            "latency_ms": 0.0,
+            "token_count": 0.0,
+            "cost_usd": 0.0,
+            "top_failures": [],
+        }
+
+    failures: list[str] = []
+    citation_support: list[float] = []
+    unsupported: list[float] = []
+    retrieval_misses = 0
+    reliability_scores: list[float] = []
+    latency: list[float] = []
+    tokens: list[float] = []
+    for trace in rows:
+        evaluation = trace.get("evaluation") or {}
+        failure = evaluation.get("failure") or {}
+        scores = evaluation.get("scores") or {}
+        failure_type = failure.get("failure_type") or "unknown"
+        if failure_type != "no_failure_detected":
+            failures.append(failure_type)
+        if failure_type == "retrieval_miss":
+            retrieval_misses += 1
+        citation_support.append(float(scores.get("citation_support") or 0.0))
+        unsupported.append(float(scores.get("unsupported_claim_rate") or 0.0))
+        reliability_scores.append(float((evaluation.get("reliability") or {}).get("score") or 0.0))
+        answer = trace.get("answer") or {}
+        usage = answer.get("usage") or {}
+        metadata = answer.get("metadata") or {}
+        if "latency_ms" in metadata:
+            latency.append(float(metadata["latency_ms"]))
+        if "total_tokens" in usage:
+            tokens.append(float(usage["total_tokens"]))
+    avg_tokens = _avg(tokens)
+    return {
+        "questions_tested": count,
+        "reliability_score": _avg(reliability_scores),
+        "failure_rate": round(len(failures) / count, 3),
+        "citation_support": _avg(citation_support),
+        "avg_citation_support": _avg(citation_support),
+        "unsupported_claim_rate": _avg(unsupported),
+        "retrieval_miss_rate": round(retrieval_misses / count, 3),
+        "latency_ms": _avg(latency),
+        "token_count": avg_tokens,
+        "cost_usd": round(avg_tokens * 0.000001, 6),
+        "top_failures": _top_failures(failures),
+    }
+
+
+def _run_question(
+    *,
+    contexttrace: ContextTrace,
+    dataset_name: str,
+    question: dict[str, Any],
+    expected_answer: str,
+    expected_sources: list[str],
+    chunks: list[dict[str, Any]],
+    strategy: str,
+    position: int,
+) -> str:
+    query = str(question["query"])
+    retrieved = _retrieve(query, chunks, top_k=STRATEGY_TOP_K.get(strategy, 3))
+    expected_failure = question.get("expected_failure")
+    if expected_sources and expected_failure in {None, "citation_mismatch"}:
+        retrieved = _force_expected_sources(chunks, expected_sources, retrieved)
+        retrieved = [
+            chunk
+            for chunk in retrieved
+            if (chunk["source"] in expected_sources or (chunk.get("metadata") or {}).get("stance") != "archived")
+        ]
+    if expected_failure == "retrieval_miss":
+        retrieved = [chunk for chunk in retrieved if chunk["source"] not in expected_sources]
+        if not retrieved:
+            retrieved = [chunk for chunk in chunks if chunk["source"] not in expected_sources][:2]
+    if expected_failure == "conflicting_sources":
+        retrieved = _force_expected_sources(chunks, expected_sources, retrieved)
+    selected = retrieved[: max(1, min(len(retrieved), STRATEGY_TOP_K.get(strategy, 3)))]
+    answer = _demo_answer(question, expected_answer)
+    citations = _demo_citations(question, answer, selected, chunks, expected_sources)
+    latency_ms = 35 + (position * 7) + len(selected) * 12
+    token_count = 80 + len(answer.split()) + sum(len(chunk["content"].split()) for chunk in selected)
+
+    with contexttrace.trace(
+        query=query,
+        metadata={
+            "dataset": dataset_name,
+            "question_id": question["id"],
+            "question_type": question.get("type"),
+            "expected_sources": expected_sources,
+            "expected_failure": expected_failure or "no_failure_detected",
+            "strategy": strategy,
+        },
+    ) as trace:
+        trace.log_retrieval(retrieved, metadata={"strategy": strategy})
+        trace.log_context(selected)
+        trace.log_answer(
+            answer,
+            model="contexttrace-demo-rag",
+            usage={"total_tokens": token_count},
+            metadata={"latency_ms": latency_ms, "strategy": strategy},
+        )
+        if citations:
+            trace.log_citations(citations)
+        trace.evaluate()
+        return str(trace.trace_id)
+
+
+def _document_chunks(documents: dict[str, str]) -> list[dict[str, Any]]:
+    chunks: list[dict[str, Any]] = []
+    for source, text in documents.items():
+        sections = [section.strip() for section in re.split(r"\n##\s+", text) if section.strip()]
+        for index, section in enumerate(sections):
+            content = re.sub(r"^#\s+", "", section).strip()
+            if not content:
+                continue
+            if len(content.split()) <= 4 and Path(source).stem.replace("_", " ").lower() in content.lower():
+                continue
+            chunks.append(
+                {
+                    "chunk_id": "%s_%s" % (Path(source).stem, index + 1),
+                    "content": content,
+                    "source": source,
+                    "metadata": {
+                        "section": content.splitlines()[0][:80],
+                        "stance": _stance(source, content),
+                    },
+                    "relevance_score": 0.0,
+                }
+            )
+    return chunks
+
+
+def _retrieve(query: str, chunks: list[dict[str, Any]], *, top_k: int) -> list[dict[str, Any]]:
+    query_terms = _terms(query)
+    scored = []
+    for chunk in chunks:
+        score = len(query_terms & _terms(chunk["content"])) / max(len(query_terms), 1)
+        scored.append({**chunk, "relevance_score": round(score, 3)})
+    return sorted(scored, key=lambda chunk: (-float(chunk["relevance_score"]), chunk["source"]))[:top_k]
+
+
+def _force_expected_sources(
+    chunks: list[dict[str, Any]],
+    expected_sources: list[str],
+    retrieved: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    forced = [chunk for chunk in chunks if chunk["source"] in expected_sources]
+    seen = {chunk["chunk_id"] for chunk in forced}
+    return forced + [chunk for chunk in retrieved if chunk["chunk_id"] not in seen]
+
+
+def _demo_answer(question: dict[str, Any], expected_answer: str) -> str:
+    failure = question.get("expected_failure")
+    if failure == "should_have_abstained":
+        return "Yes, the documents say this exception is available to eligible employees or customers."
+    if failure == "unsupported_answer":
+        return "Yes, the policy gives a special exception that is not limited by the documented rules."
+    return expected_answer
+
+
+def _demo_citations(
+    question: dict[str, Any],
+    answer: str,
+    selected: list[dict[str, Any]],
+    chunks: list[dict[str, Any]],
+    expected_sources: list[str],
+) -> list[dict[str, Any]]:
+    if not answer or not selected:
+        return []
+    failure = question.get("expected_failure")
+    claim = answer.split(".")[0].strip() + "."
+    if failure == "citation_mismatch":
+        wrong = next((chunk for chunk in selected if chunk["source"] not in expected_sources), None)
+        if wrong is None:
+            wrong = next((chunk for chunk in chunks if chunk["source"] not in expected_sources), selected[0])
+        return [{"claim": claim, "source_chunk_id": wrong["chunk_id"]}]
+    return [{"claim": claim, "source_chunk_id": selected[0]["chunk_id"]}]
+
+
+def _stance(source: str, content: str) -> str:
+    lowered = (source + " " + content).lower()
+    if "archived" in lowered or "old_" in lowered or "legacy" in lowered:
+        return "archived"
+    if "current" in lowered or "policy" in lowered:
+        return "current"
+    return "neutral"
+
+
+def _terms(text: str) -> set[str]:
+    return {
+        token.strip(".,:;!?()[]{}\"'").lower().rstrip("s")
+        for token in text.split()
+        if len(token.strip(".,:;!?()[]{}\"'")) > 2
+    }
+
+
+def _avg(values: list[float]) -> float:
+    return round(sum(values) / len(values), 3) if values else 0.0
+
+
+def _top_failures(failures: list[str]) -> list[str]:
+    counts = {failure: failures.count(failure) for failure in set(failures)}
+    return [name for name, _ in sorted(counts.items(), key=lambda item: (-item[1], item[0]))[:5]]