Spaces:
Sleeping
Sleeping
File size: 1,819 Bytes
6a8cd1b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
"""
Data models for Context Thread Agent
"""
from typing import List, Dict, Optional, Any
from dataclasses import dataclass
from enum import Enum
class CellType(Enum):
CODE = "code"
MARKDOWN = "markdown"
RAW = "raw"
@dataclass
class Cell:
"""Represents a Jupyter notebook cell."""
cell_id: str
cell_type: CellType
source: str
outputs: List[Dict[str, Any]] = None
def __post_init__(self):
if self.outputs is None:
self.outputs = []
@dataclass
class ContextUnit:
"""A cell with its context and dependencies."""
cell: Cell
intent: str
dependencies: List[str]
context_window: List[str] = None
def __post_init__(self):
if self.context_window is None:
self.context_window = []
@dataclass
class ContextThread:
"""A thread of related context units."""
notebook_name: str
thread_id: str
units: List[ContextUnit]
metadata: Dict[str, Any] = None
def __post_init__(self):
if self.metadata is None:
self.metadata = {}
@dataclass
class QueryRequest:
"""A user query request."""
query: str
notebook_path: Optional[str] = None
top_k: int = 5
@dataclass
class Citation:
"""A citation to a specific cell."""
cell_id: str
cell_type: CellType
content_snippet: str
intent: Optional[str] = None
@dataclass
class AgentResponse:
"""Response from the agent."""
answer: str
citations: List[Citation]
confidence: float
has_hallucination_risk: bool
retrieved_units: List[ContextUnit]
@dataclass
class RetrievalResult:
"""Result from retrieval system."""
units: List[ContextUnit]
scores: List[float]
query: str |