rag-debugger 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,70 @@
1
+ # IDE
2
+ .idea
3
+ *.iml
4
+ .vscode/settings.json
5
+
6
+ # misc
7
+ .DS_Store
8
+ npm-debug*
9
+ **/*.tsbuildinfo
10
+
11
+ # Node
12
+ **/node_modules/
13
+ **/.pnpm-store/
14
+
15
+ # Next.js
16
+ **/.next/
17
+ **/out/
18
+
19
+ # Build
20
+ **/build/
21
+ **/dist/
22
+
23
+ # Python
24
+ **/__pycache__/
25
+ **/*.pyc
26
+ **/*.pyo
27
+ **/.venv/
28
+ **/venv/
29
+ *.egg-info/
30
+ **/*.egg-info/
31
+ **/.pytest_cache/
32
+
33
+ # DuckDB
34
+ *.duckdb
35
+ *.duckdb.wal
36
+
37
+ # Env
38
+ *.swp
39
+ *.env
40
+ *.env.*
41
+ *.env.local
42
+
43
+ # Certs
44
+ *.pem
45
+
46
+ # Turbo
47
+ .turbo
48
+
49
+ # Firebase
50
+ **/.firebase/
51
+
52
+ # Configs
53
+ **/public/conf.js
54
+
55
+ # Archives
56
+ *.zip
57
+
58
+ # Testing
59
+ **/test-results/
60
+ **/playwright-report/
61
+ **/blob-report/
62
+ **/playwright/.cache/
63
+
64
+ # Million Lint
65
+ .million
66
+
67
+ # Sentence-transformers cache
68
+ **/sentence_transformers/
69
+ **/*.duckdb
70
+ **/*.duckdb.wal
@@ -0,0 +1,174 @@
1
+ Metadata-Version: 2.4
2
+ Name: rag-debugger
3
+ Version: 1.0.0
4
+ Summary: Real-time debugging SDK for RAG pipelines
5
+ Project-URL: Homepage, https://github.com/ChanduBobbili/rag-debugger
6
+ Project-URL: Repository, https://github.com/ChanduBobbili/rag-debugger
7
+ Project-URL: Issues, https://github.com/ChanduBobbili/rag-debugger/issues
8
+ Author: Chandu Bobbili
9
+ License: MIT
10
+ Keywords: debugging,llm,observability,rag,tracing
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: License :: OSI Approved :: MIT License
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Software Development :: Debuggers
20
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
21
+ Requires-Python: >=3.10
22
+ Requires-Dist: httpx>=0.24.0
23
+ Requires-Dist: pydantic>=2.0.0
24
+ Provides-Extra: all
25
+ Requires-Dist: langchain-core>=0.1.0; extra == 'all'
26
+ Requires-Dist: llama-index-core>=0.10.0; extra == 'all'
27
+ Requires-Dist: openai>=1.0.0; extra == 'all'
28
+ Provides-Extra: langchain
29
+ Requires-Dist: langchain-core>=0.1.0; extra == 'langchain'
30
+ Provides-Extra: llamaindex
31
+ Requires-Dist: llama-index-core>=0.10.0; extra == 'llamaindex'
32
+ Provides-Extra: openai
33
+ Requires-Dist: openai>=1.0.0; extra == 'openai'
34
+ Description-Content-Type: text/markdown
35
+
36
+ # RAG Debugger SDK 🔍
37
+
38
+ [![PyPI version](https://img.shields.io/pypi/v/rag-debugger.svg)](https://pypi.org/project/rag-debugger/)
39
+ [![Python](https://img.shields.io/pypi/pyversions/rag-debugger.svg)](https://pypi.org/project/rag-debugger/)
40
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
41
+
42
+ **One-line decorator to debug your RAG pipelines in real time.**
43
+
44
+ Instrument any Python RAG pipeline with `@rag_trace` — captures inputs, outputs, timing, and errors for every stage (embed → retrieve → rerank → generate) and streams them to the [RAG Debugger Dashboard](https://github.com/ChanduBobbili/rag-debugger).
45
+
46
+ ## Features
47
+
48
+ - 🔗 **One decorator** — `@rag_trace("retrieve")` on your existing functions
49
+ - ⚡ **Non-blocking** — async background worker, never slows your pipeline
50
+ - 🧵 **Auto-correlation** — `trace_id` / `query_id` via `ContextVar` (no manual threading)
51
+ - 🔒 **PII scrubbing** — emails, phone numbers, SSNs, API keys automatically redacted
52
+ - 🔌 **Framework adapters** — LangChain, LlamaIndex, and OpenAI out of the box
53
+ - 🛡️ **Safe** — errors in the SDK never crash your application
54
+
55
+ ## Installation
56
+
57
+ ```bash
58
+ pip install rag-debugger
59
+ ```
60
+
61
+ With framework adapters:
62
+
63
+ ```bash
64
+ pip install rag-debugger[langchain] # LangChain
65
+ pip install rag-debugger[llamaindex] # LlamaIndex
66
+ pip install rag-debugger[openai] # OpenAI
67
+ pip install rag-debugger[all] # All adapters
68
+ ```
69
+
70
+ ## Quick Start
71
+
72
+ ```python
73
+ from rag_debugger import init, rag_trace
74
+
75
+ # 1. Point to your RAG Debugger server
76
+ init(dashboard_url="http://localhost:7777")
77
+
78
+ # 2. Decorate your pipeline functions
79
+ @rag_trace("embed")
80
+ async def embed_query(query: str) -> list[float]:
81
+ return await my_embedder.embed(query)
82
+
83
+ @rag_trace("retrieve")
84
+ async def retrieve_chunks(vector: list[float], k: int = 10):
85
+ return await vector_store.query(vector, k)
86
+
87
+ @rag_trace("rerank")
88
+ async def rerank(query: str, chunks: list) -> list:
89
+ return await reranker.rerank(query, chunks)
90
+
91
+ @rag_trace("generate")
92
+ async def generate(query: str, context: str) -> str:
93
+ return await llm.complete(query, context)
94
+
95
+ # 3. Call your pipeline — traces appear in the dashboard
96
+ answer = await generate(query, context)
97
+ ```
98
+
99
+ The decorator automatically:
100
+
101
+ - Generates `trace_id` and `query_id` per request
102
+ - Captures function inputs and outputs
103
+ - Measures `duration_ms` for each stage
104
+ - Emits a `session_complete` summary after the generate stage
105
+ - Scrubs PII before sending
106
+
107
+ ## Framework Adapters
108
+
109
+ ### LangChain
110
+
111
+ ```python
112
+ from rag_debugger.adapters.langchain import RAGDebuggerCallback
113
+
114
+ handler = RAGDebuggerCallback()
115
+ chain.invoke({"query": "..."}, config={"callbacks": [handler]})
116
+ ```
117
+
118
+ ### LlamaIndex
119
+
120
+ ```python
121
+ from rag_debugger.adapters.llamaindex import RAGDebuggerLlamaIndex
122
+ from llama_index.core.callbacks import CallbackManager
123
+
124
+ handler = RAGDebuggerLlamaIndex()
125
+ callback_manager = CallbackManager([handler])
126
+ index = VectorStoreIndex.from_documents(docs, callback_manager=callback_manager)
127
+ ```
128
+
129
+ ### OpenAI
130
+
131
+ ```python
132
+ from rag_debugger.adapters.openai import RAGDebuggerOpenAI
133
+
134
+ client = RAGDebuggerOpenAI(openai.AsyncOpenAI())
135
+ embedding = await client.embed("What is RAG?")
136
+ response = await client.complete("Explain RAG")
137
+ ```
138
+
139
+ ## Advanced Usage
140
+
141
+ ### Explicit Trace Control
142
+
143
+ ```python
144
+ from rag_debugger import new_trace, reset_context
145
+
146
+ # Group events under a custom trace
147
+ new_trace(trace_id="my-trace-123", query_id="q-001")
148
+ await embed_query("What is RAG?")
149
+ await retrieve_chunks(vector)
150
+
151
+ # Reset for the next request
152
+ reset_context()
153
+ ```
154
+
155
+ ### Async Context Manager
156
+
157
+ ```python
158
+ import rag_debugger
159
+
160
+ async with rag_debugger.trace(trace_id="req-123") as t:
161
+ print(t.trace_id)
162
+ result = await my_rag_pipeline(query)
163
+ # Context is automatically restored after the block
164
+ ```
165
+
166
+ ## Documentation
167
+
168
+ - [Full SDK Documentation](https://github.com/ChanduBobbili/rag-debugger/blob/main/docs/SDK.md)
169
+ - [Server Documentation](https://github.com/ChanduBobbili/rag-debugger/blob/main/docs/SERVER.md)
170
+ - [Dashboard Documentation](https://github.com/ChanduBobbili/rag-debugger/blob/main/docs/DASHBOARD.md)
171
+
172
+ ## License
173
+
174
+ MIT — see [LICENSE](https://github.com/ChanduBobbili/rag-debugger/blob/main/LICENSE) for details.
@@ -0,0 +1,139 @@
1
+ # RAG Debugger SDK 🔍
2
+
3
+ [![PyPI version](https://img.shields.io/pypi/v/rag-debugger.svg)](https://pypi.org/project/rag-debugger/)
4
+ [![Python](https://img.shields.io/pypi/pyversions/rag-debugger.svg)](https://pypi.org/project/rag-debugger/)
5
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
6
+
7
+ **One-line decorator to debug your RAG pipelines in real time.**
8
+
9
+ Instrument any Python RAG pipeline with `@rag_trace` — captures inputs, outputs, timing, and errors for every stage (embed → retrieve → rerank → generate) and streams them to the [RAG Debugger Dashboard](https://github.com/ChanduBobbili/rag-debugger).
10
+
11
+ ## Features
12
+
13
+ - 🔗 **One decorator** — `@rag_trace("retrieve")` on your existing functions
14
+ - ⚡ **Non-blocking** — async background worker, never slows your pipeline
15
+ - 🧵 **Auto-correlation** — `trace_id` / `query_id` via `ContextVar` (no manual threading)
16
+ - 🔒 **PII scrubbing** — emails, phone numbers, SSNs, API keys automatically redacted
17
+ - 🔌 **Framework adapters** — LangChain, LlamaIndex, and OpenAI out of the box
18
+ - 🛡️ **Safe** — errors in the SDK never crash your application
19
+
20
+ ## Installation
21
+
22
+ ```bash
23
+ pip install rag-debugger
24
+ ```
25
+
26
+ With framework adapters:
27
+
28
+ ```bash
29
+ pip install rag-debugger[langchain] # LangChain
30
+ pip install rag-debugger[llamaindex] # LlamaIndex
31
+ pip install rag-debugger[openai] # OpenAI
32
+ pip install rag-debugger[all] # All adapters
33
+ ```
34
+
35
+ ## Quick Start
36
+
37
+ ```python
38
+ from rag_debugger import init, rag_trace
39
+
40
+ # 1. Point to your RAG Debugger server
41
+ init(dashboard_url="http://localhost:7777")
42
+
43
+ # 2. Decorate your pipeline functions
44
+ @rag_trace("embed")
45
+ async def embed_query(query: str) -> list[float]:
46
+ return await my_embedder.embed(query)
47
+
48
+ @rag_trace("retrieve")
49
+ async def retrieve_chunks(vector: list[float], k: int = 10):
50
+ return await vector_store.query(vector, k)
51
+
52
+ @rag_trace("rerank")
53
+ async def rerank(query: str, chunks: list) -> list:
54
+ return await reranker.rerank(query, chunks)
55
+
56
+ @rag_trace("generate")
57
+ async def generate(query: str, context: str) -> str:
58
+ return await llm.complete(query, context)
59
+
60
+ # 3. Call your pipeline — traces appear in the dashboard
61
+ answer = await generate(query, context)
62
+ ```
63
+
64
+ The decorator automatically:
65
+
66
+ - Generates `trace_id` and `query_id` per request
67
+ - Captures function inputs and outputs
68
+ - Measures `duration_ms` for each stage
69
+ - Emits a `session_complete` summary after the generate stage
70
+ - Scrubs PII before sending
71
+
72
+ ## Framework Adapters
73
+
74
+ ### LangChain
75
+
76
+ ```python
77
+ from rag_debugger.adapters.langchain import RAGDebuggerCallback
78
+
79
+ handler = RAGDebuggerCallback()
80
+ chain.invoke({"query": "..."}, config={"callbacks": [handler]})
81
+ ```
82
+
83
+ ### LlamaIndex
84
+
85
+ ```python
86
+ from rag_debugger.adapters.llamaindex import RAGDebuggerLlamaIndex
87
+ from llama_index.core.callbacks import CallbackManager
88
+
89
+ handler = RAGDebuggerLlamaIndex()
90
+ callback_manager = CallbackManager([handler])
91
+ index = VectorStoreIndex.from_documents(docs, callback_manager=callback_manager)
92
+ ```
93
+
94
+ ### OpenAI
95
+
96
+ ```python
97
+ from rag_debugger.adapters.openai import RAGDebuggerOpenAI
98
+
99
+ client = RAGDebuggerOpenAI(openai.AsyncOpenAI())
100
+ embedding = await client.embed("What is RAG?")
101
+ response = await client.complete("Explain RAG")
102
+ ```
103
+
104
+ ## Advanced Usage
105
+
106
+ ### Explicit Trace Control
107
+
108
+ ```python
109
+ from rag_debugger import new_trace, reset_context
110
+
111
+ # Group events under a custom trace
112
+ new_trace(trace_id="my-trace-123", query_id="q-001")
113
+ await embed_query("What is RAG?")
114
+ await retrieve_chunks(vector)
115
+
116
+ # Reset for the next request
117
+ reset_context()
118
+ ```
119
+
120
+ ### Async Context Manager
121
+
122
+ ```python
123
+ import rag_debugger
124
+
125
+ async with rag_debugger.trace(trace_id="req-123") as t:
126
+ print(t.trace_id)
127
+ result = await my_rag_pipeline(query)
128
+ # Context is automatically restored after the block
129
+ ```
130
+
131
+ ## Documentation
132
+
133
+ - [Full SDK Documentation](https://github.com/ChanduBobbili/rag-debugger/blob/main/docs/SDK.md)
134
+ - [Server Documentation](https://github.com/ChanduBobbili/rag-debugger/blob/main/docs/SERVER.md)
135
+ - [Dashboard Documentation](https://github.com/ChanduBobbili/rag-debugger/blob/main/docs/DASHBOARD.md)
136
+
137
+ ## License
138
+
139
+ MIT — see [LICENSE](https://github.com/ChanduBobbili/rag-debugger/blob/main/LICENSE) for details.
@@ -0,0 +1,45 @@
1
+ [project]
2
+ name = "rag-debugger"
3
+ version = "1.0.0"
4
+ description = "Real-time debugging SDK for RAG pipelines"
5
+ readme = "README.md"
6
+ license = {text = "MIT"}
7
+ requires-python = ">=3.10"
8
+ authors = [
9
+ {name = "Chandu Bobbili"},
10
+ ]
11
+ keywords = ["rag", "debugging", "llm", "observability", "tracing"]
12
+ classifiers = [
13
+ "Development Status :: 3 - Alpha",
14
+ "Intended Audience :: Developers",
15
+ "License :: OSI Approved :: MIT License",
16
+ "Programming Language :: Python :: 3",
17
+ "Programming Language :: Python :: 3.10",
18
+ "Programming Language :: Python :: 3.11",
19
+ "Programming Language :: Python :: 3.12",
20
+ "Programming Language :: Python :: 3.13",
21
+ "Topic :: Software Development :: Libraries :: Python Modules",
22
+ "Topic :: Software Development :: Debuggers",
23
+ ]
24
+ dependencies = [
25
+ "httpx>=0.24.0",
26
+ "pydantic>=2.0.0",
27
+ ]
28
+
29
+ [project.optional-dependencies]
30
+ langchain = ["langchain-core>=0.1.0"]
31
+ llamaindex = ["llama-index-core>=0.10.0"]
32
+ openai = ["openai>=1.0.0"]
33
+ all = ["langchain-core>=0.1.0", "llama-index-core>=0.10.0", "openai>=1.0.0"]
34
+
35
+ [project.urls]
36
+ Homepage = "https://github.com/ChanduBobbili/rag-debugger"
37
+ Repository = "https://github.com/ChanduBobbili/rag-debugger"
38
+ Issues = "https://github.com/ChanduBobbili/rag-debugger/issues"
39
+
40
+ [tool.hatch.build.targets.wheel]
41
+ packages = ["rag_debugger"]
42
+
43
+ [build-system]
44
+ requires = ["hatchling"]
45
+ build-backend = "hatchling.build"
@@ -0,0 +1,77 @@
1
+ import uuid
2
+ from contextlib import asynccontextmanager
3
+ from .emitter import configure, stop_worker
4
+ from .decorators import rag_trace
5
+ from .context import set_trace_id, set_query_id, reset_context
6
+ from .context import _trace_id, _query_id
7
+
8
+ __version__ = "1.0.0"
9
+
10
+ _initialized = False
11
+
12
+
13
+ def init(dashboard_url: str = "http://localhost:7777") -> None:
14
+ """Call once at application startup.
15
+
16
+ Configures the dashboard URL. The background worker starts lazily
17
+ on the first ``emit()`` call, so this is safe to call at import time
18
+ or before the async event loop is running.
19
+ """
20
+ global _initialized
21
+ configure(dashboard_url)
22
+ _initialized = True
23
+
24
+
25
+ def new_trace(
26
+ trace_id: str | None = None,
27
+ query_id: str | None = None,
28
+ ) -> None:
29
+ """Explicitly set trace/query IDs (optional — auto-generated if not called)."""
30
+ if trace_id:
31
+ set_trace_id(trace_id)
32
+ if query_id:
33
+ set_query_id(query_id)
34
+
35
+
36
+ class _TraceHandle:
37
+ """Lightweight handle returned by the ``trace()`` context manager."""
38
+ __slots__ = ("trace_id", "query_id")
39
+
40
+ def __init__(self, trace_id: str, query_id: str) -> None:
41
+ self.trace_id = trace_id
42
+ self.query_id = query_id
43
+
44
+
45
+ @asynccontextmanager
46
+ async def trace(
47
+ trace_id: str | None = None,
48
+ query_id: str | None = None,
49
+ ):
50
+ """Async context manager for explicit trace scoping.
51
+
52
+ Usage::
53
+
54
+ async with rag_debugger.trace(trace_id="req-123") as t:
55
+ print(t.trace_id)
56
+ result = await my_rag_pipeline(query)
57
+ # Context is automatically restored after the block
58
+
59
+ Nested ``trace()`` contexts work correctly — the outer context
60
+ is restored when the inner block exits.
61
+ """
62
+ tid = trace_id or str(uuid.uuid4())
63
+ qid = query_id or str(uuid.uuid4())
64
+
65
+ # Save previous values using ContextVar tokens
66
+ trace_token = _trace_id.set(tid)
67
+ query_token = _query_id.set(qid)
68
+
69
+ try:
70
+ yield _TraceHandle(tid, qid)
71
+ finally:
72
+ # Restore previous values
73
+ _trace_id.reset(trace_token)
74
+ _query_id.reset(query_token)
75
+
76
+
77
+ __all__ = ["init", "rag_trace", "new_trace", "reset_context", "trace", "stop_worker", "__version__"]
File without changes
@@ -0,0 +1,80 @@
1
+ try:
2
+ from langchain_core.callbacks import BaseCallbackHandler
3
+ from langchain_core.outputs import LLMResult
4
+ except ImportError:
5
+ raise ImportError(
6
+ "LangChain adapter requires langchain-core. "
7
+ "Install with: pip install rag-debugger[langchain]"
8
+ )
9
+
10
+ import asyncio
11
+ import time
12
+ import uuid
13
+ from ..context import get_or_create_trace_id, get_or_create_query_id
14
+ from ..emitter import emit
15
+
16
+
17
+ class RAGDebuggerCallback(BaseCallbackHandler):
18
+ """
19
+ LangChain callback handler.
20
+ Usage:
21
+ from rag_debugger.adapters.langchain import RAGDebuggerCallback
22
+ handler = RAGDebuggerCallback()
23
+ chain.invoke({"query": "..."}, config={"callbacks": [handler]})
24
+ """
25
+
26
+ def __init__(self) -> None:
27
+ self._retriever_start: float = 0
28
+ self._llm_start: float = 0
29
+ self._query_text: str = ""
30
+
31
+ def on_retriever_start(self, serialized, query, **kwargs) -> None:
32
+ self._retriever_start = time.time()
33
+ self._query_text = query
34
+
35
+ def on_retriever_end(self, documents, **kwargs) -> None:
36
+ duration = (time.time() - self._retriever_start) * 1000
37
+ chunks = [
38
+ {
39
+ "chunk_id": str(i),
40
+ "text": doc.page_content[:1000],
41
+ "cosine_score": doc.metadata.get("score", 0.0),
42
+ "final_rank": i,
43
+ "metadata": doc.metadata,
44
+ }
45
+ for i, doc in enumerate(documents)
46
+ ]
47
+ try:
48
+ loop = asyncio.get_running_loop()
49
+ loop.create_task(emit({
50
+ "id": str(uuid.uuid4()),
51
+ "trace_id": get_or_create_trace_id(),
52
+ "query_id": get_or_create_query_id(),
53
+ "stage": "retrieve",
54
+ "ts_start": self._retriever_start,
55
+ "duration_ms": duration,
56
+ "query_text": self._query_text,
57
+ "chunks": chunks,
58
+ }))
59
+ except RuntimeError:
60
+ pass # No running loop — skip
61
+
62
+ def on_llm_start(self, serialized, prompts, **kwargs) -> None:
63
+ self._llm_start = time.time()
64
+
65
+ def on_llm_end(self, response: LLMResult, **kwargs) -> None:
66
+ duration = (time.time() - self._llm_start) * 1000
67
+ answer = response.generations[0][0].text if response.generations else ""
68
+ try:
69
+ loop = asyncio.get_running_loop()
70
+ loop.create_task(emit({
71
+ "id": str(uuid.uuid4()),
72
+ "trace_id": get_or_create_trace_id(),
73
+ "query_id": get_or_create_query_id(),
74
+ "stage": "generate",
75
+ "ts_start": self._llm_start,
76
+ "duration_ms": duration,
77
+ "generated_answer": answer,
78
+ }))
79
+ except RuntimeError:
80
+ pass
@@ -0,0 +1,105 @@
1
+ """LlamaIndex observer adapter for RAG Debugger SDK."""
2
+
3
+ try:
4
+ from llama_index.core.callbacks import CallbackManager, CBEventType, LlamaDebugHandler
5
+ from llama_index.core.callbacks.base_handler import BaseCallbackHandler
6
+ except ImportError:
7
+ raise ImportError(
8
+ "LlamaIndex adapter requires llama-index-core. "
9
+ "Install with: pip install rag-debugger[llamaindex]"
10
+ )
11
+
12
+ import asyncio
13
+ import time
14
+ import uuid
15
+ from typing import Any, Dict, List, Optional
16
+ from ..context import get_or_create_trace_id, get_or_create_query_id
17
+ from ..emitter import emit
18
+
19
+
20
+ class RAGDebuggerLlamaIndex(BaseCallbackHandler):
21
+ """
22
+ LlamaIndex callback handler for RAG Debugger.
23
+ Usage:
24
+ from rag_debugger.adapters.llamaindex import RAGDebuggerLlamaIndex
25
+ handler = RAGDebuggerLlamaIndex()
26
+ callback_manager = CallbackManager([handler])
27
+ index = VectorStoreIndex.from_documents(docs, callback_manager=callback_manager)
28
+ """
29
+
30
+ def __init__(self) -> None:
31
+ super().__init__([], [])
32
+ self._event_starts: Dict[str, float] = {}
33
+
34
+ def on_event_start(
35
+ self,
36
+ event_type: CBEventType,
37
+ payload: Optional[Dict[str, Any]] = None,
38
+ event_id: str = "",
39
+ **kwargs,
40
+ ) -> str:
41
+ self._event_starts[event_id] = time.time()
42
+ return event_id
43
+
44
+ def on_event_end(
45
+ self,
46
+ event_type: CBEventType,
47
+ payload: Optional[Dict[str, Any]] = None,
48
+ event_id: str = "",
49
+ **kwargs,
50
+ ) -> None:
51
+ start_time = self._event_starts.pop(event_id, time.time())
52
+ duration = (time.time() - start_time) * 1000
53
+
54
+ stage = self._map_event_type(event_type)
55
+ if stage is None:
56
+ return
57
+
58
+ event = {
59
+ "id": str(uuid.uuid4()),
60
+ "trace_id": get_or_create_trace_id(),
61
+ "query_id": get_or_create_query_id(),
62
+ "stage": stage,
63
+ "ts_start": start_time,
64
+ "duration_ms": duration,
65
+ }
66
+
67
+ if payload:
68
+ if stage == "retrieve" and "nodes" in payload:
69
+ event["chunks"] = [
70
+ {
71
+ "chunk_id": str(i),
72
+ "text": str(getattr(n, "text", ""))[:1000],
73
+ "cosine_score": float(getattr(n, "score", 0.0)),
74
+ "final_rank": i,
75
+ }
76
+ for i, n in enumerate(payload["nodes"])
77
+ ]
78
+ elif stage == "generate" and "response" in payload:
79
+ event["generated_answer"] = str(payload["response"])
80
+
81
+ try:
82
+ loop = asyncio.get_running_loop()
83
+ loop.create_task(emit(event))
84
+ except RuntimeError:
85
+ pass
86
+
87
+ def start_trace(self, trace_id: Optional[str] = None) -> None:
88
+ pass
89
+
90
+ def end_trace(
91
+ self,
92
+ trace_id: Optional[str] = None,
93
+ trace_map: Optional[Dict[str, List[str]]] = None,
94
+ ) -> None:
95
+ pass
96
+
97
+ @staticmethod
98
+ def _map_event_type(event_type: CBEventType) -> Optional[str]:
99
+ mapping = {
100
+ CBEventType.EMBEDDING: "embed",
101
+ CBEventType.RETRIEVE: "retrieve",
102
+ CBEventType.RERANKING: "rerank",
103
+ CBEventType.LLM: "generate",
104
+ }
105
+ return mapping.get(event_type)