上下文与决策智能

86@dataclass87class Decision:88    """Core decision data model with full context tracking."""89    90    decision_id: str91    category: str92    scenario: str93    reasoning: str94    outcome: str95    confidence: float96    timestamp: datetime97    decision_maker: str98    reasoning_embedding: Optional[List[float]] = None99    node2vec_embedding: Optional[List[float]] = None100    valid_from: Optional[str] = None101    valid_until: Optional[str] = None102    metadata: Dict[str, Any] = field(default_factory=dict)103    104    def __post_init__(self, auto_generate_id: bool = True):105        """Validate decision data."""106        if auto_generate_id and not self.decision_id:  # Handle both None and empty string107            self.decision_id = str(uuid.uuid4())108        elif not self.decision_id and not auto_generate_id:109            raise ValueError("decision_id is required when auto_generate_id=False")110        if not 0 <= self.confidence <= 1:111            raise ValueError("Confidence must be between 0 and 1")112    113    def to_dict(self) -> Dict[str, Any]:114        """Convert decision to dictionary."""115        return {116            "decision_id": self.decision_id,117            "category": self.category,118            "scenario": self.scenario,119            "reasoning": self.reasoning,120            "outcome": self.outcome,121            "confidence": self.confidence,122            "timestamp": self.timestamp.isoformat(),123            "decision_maker": self.decision_maker,124            "reasoning_embedding": self.reasoning_embedding,125            "node2vec_embedding": self.node2vec_embedding,126            "valid_from": self.valid_from,127            "valid_until": self.valid_until,128            "metadata": self.metadata129        }130    131    @classmethod132    def from_dict(cls, data: Dict[str, Any]) -> "Decision":133        """Create decision from dictionary."""134        if isinstance(data.get("timestamp"), str):135            data["timestamp"] = datetime.fromisoformat(data["timestamp"])136        return cls(**data)

321@dataclass322class ContextNode:323    """Context graph node (Internal implementation)."""324325    node_id: str326    node_type: str327    content: str328    metadata: Dict[str, Any] = field(default_factory=dict)329    properties: Dict[str, Any] = field(default_factory=dict)330    valid_from: Optional[str] = None  331    valid_until: Optional[str] = None  332333    def is_active(self, at_time: Optional[datetime] = None) -> bool:334        """Return True if this node is active at the given time (defaults to now).335336        Both ``at_time`` and stored bounds are normalized to tz-naive UTC so that337        callers may pass either aware or naive datetimes without raising TypeError.338        """339        if self.valid_from is None and self.valid_until is None:340            return True341        now = at_time if at_time is not None else datetime.utcnow()342        if now.tzinfo is not None:343            now = now.astimezone(timezone.utc).replace(tzinfo=None)344        start = _parse_iso_dt(self.valid_from) if self.valid_from is not None else None345        end = _parse_iso_dt(self.valid_until) if self.valid_until is not None else None346        if start is not None and now < start:347            return False348        if end is not None and now > end:349            return False350        return True351352    def to_dict(self) -> Dict[str, Any]:353        """Convert to dictionary format."""354        props = self.properties.copy()355        props.update(self.metadata)356        props["content"] = self.content357        if self.valid_from is not None:358            props["valid_from"] = self.valid_from359        if self.valid_until is not None:360            props["valid_until"] = self.valid_until361        return {"id": self.node_id, "type": self.node_type, "properties": props}

115    def record_decision(116        self,117        decision: Decision,118        entities: List[str],119        source_documents: List[str]120    ) -> str:121        """122        Record decision with full context.123        124        Args:125            decision: Decision object to record126            entities: List of entity IDs linked to this decision127            source_documents: List of source document identifiers128            129        Returns:130            Decision ID131        """132        try:133            # Generate embeddings if available134            if self.embedding_generator and not decision.reasoning_embedding:135                decision.reasoning_embedding = self.embedding_generator.generate(136                    decision.reasoning137                )138            139            # Store decision in graph database140            self._store_decision_node(decision)141            142            # Link to entities143            self.link_entities(decision.decision_id, entities)144            145            # Track provenance146            if self.provenance_manager:147                self._track_decision_provenance(decision, source_documents)148            149            self.logger.info(f"Recorded decision: {decision.decision_id} | Actor: {decision.decision_maker} | Timestamp: {decision.timestamp} | Outcome: {decision.outcome} | Category: {decision.category}")150            return decision.decision_id151            152        except Exception as e:153            self.logger.exception("Failed to record decision")154            raise

469    def _store_decision_node(self, decision: Decision) -> None:470        """Store decision node in graph database."""471        metadata = decision.metadata.copy() if decision.metadata else {}472        metadata.update({473            "category": decision.category,474            "scenario": decision.scenario,475            "reasoning": decision.reasoning,476            "outcome": decision.outcome,477            "confidence": decision.confidence,478            "timestamp": decision.timestamp.isoformat() if decision.timestamp else None,479            "decision_maker": decision.decision_maker,480            "reasoning_embedding": decision.reasoning_embedding,481            "node2vec_embedding": decision.node2vec_embedding482        })483        484        if type(self.graph_store) is ContextGraph:485            self.graph_store.add_node(486                node_id=decision.decision_id,487                node_type="Decision",488                **metadata489            )490            return491492        query = """493        CREATE (d:Decision {494            decision_id: $decision_id,495            category: $category,496            scenario: $scenario,497            reasoning: $reasoning,498            outcome: $outcome,499            confidence: $confidence,500            timestamp: $timestamp,501            decision_maker: $decision_maker,502            reasoning_embedding: $reasoning_embedding,503            node2vec_embedding: $node2vec_embedding,504            metadata: $metadata505        })506        """507        self.graph_store.execute_query(query, {508            "decision_id": decision.decision_id,509            "category": decision.category,510            "scenario": decision.scenario,511            "reasoning": decision.reasoning,512            "outcome": decision.outcome,513            "confidence": decision.confidence,514            "timestamp": decision.timestamp,515            "decision_maker": decision.decision_maker,516            "reasoning_embedding": decision.reasoning_embedding,517            "node2vec_embedding": decision.node2vec_embedding,518            "metadata": decision.metadata519        })

587    def _track_decision_provenance(588        self,589        decision: Decision,590        source_documents: List[str]591    ) -> None:592        """Track decision provenance using ProvenanceManager."""593        if not self.provenance_manager:594            return595        596        try:597            # Track decision as entity598            self.provenance_manager.track_entity(599                entity_id=decision.decision_id,600                entity_type="decision",601                activity_id="decision_making_process",602                agent_id=decision.decision_maker,603                source_documents=source_documents,604                confidence=decision.confidence605            )606            607            # Track decision-making activity608            self.provenance_manager.track_activity(609                activity_id=f"decision_{decision.decision_id}",610                activity_type="decision_making",611                agent_id=decision.decision_maker,612                used_entities=[decision.decision_id],613                started_at=decision.timestamp,614                ended_at=decision.timestamp615            )616            617        except Exception as e:618            self.logger.exception("Failed to track provenance")

226                query = """227                MATCH (d:Decision {decision_id: $decision_id})228                MATCH (p:Policy {policy_id: $policy_id})229                WHERE $policy_version IS NULL OR p.version = $policy_version230                WITH d, p231                ORDER BY p.updated_at DESC, p.version DESC232                LIMIT 1233                MERGE (d)-[r:APPLIED_POLICY]->(p)234                SET r.policy_id = $policy_id,235                    r.policy_version = p.version,236                    d.applied_at = timestamp()237                RETURN p.policy_id as policy_id, p.version as version238                """

2751    def add_causal_relationship(2752        self,2753        source_decision_id: str,2754        target_decision_id: str,2755        relationship_type: str2756    ) -> None:2757        """2758        Add causal relationship between decisions.2759        2760        Args:2761            source_decision_id: Source decision ID2762            target_decision_id: Target decision ID2763            relationship_type: Type of relationship (CAUSED, INFLUENCED, PRECEDENT_FOR)2764        """2765        # Normalize so callers may use either vocabulary's spelling2766        # ("causes" from CausalChainAnalyzer, or "CAUSED" from this module's2767        # canonical constant); the stored form is always canonical. Invalid2768        # inputs keep raising ValueError rather than AttributeError.2769        if not isinstance(relationship_type, str):2770            raise ValueError(f"Relationship type must be one of: {_CAUSAL_EDGE_TYPES}")2771        relationship_type = _CAUSAL_EDGE_ALIASES.get(relationship_type.strip().upper())2772        if relationship_type is None:2773            raise ValueError(f"Relationship type must be one of: {_CAUSAL_EDGE_TYPES}")2774        2775        # Check if decisions exist - if not, skip adding relationship2776        if source_decision_id not in self.nodes or target_decision_id not in self.nodes:2777            return2778        2779        # Check if nodes are decision nodes - if not, skip adding relationship2780        source_node = self.nodes[source_decision_id]2781        target_node = self.nodes[target_decision_id]2782        if (not hasattr(source_node, 'node_type') or not isinstance(source_node.node_type, str) or2783            not hasattr(target_node, 'node_type') or not isinstance(target_node.node_type, str) or2784            source_node.node_type.lower() != "decision" or 2785            target_node.node_type.lower() != "decision"):2786            return2787        2788        edge = ContextEdge(2789            source_id=source_decision_id,2790            target_id=target_decision_id,2791            edge_type=relationship_type,2792            weight=1.0,2793            metadata={"recorded_at": datetime.utcnow().isoformat()},2794        )2795        self._add_internal_edge(edge)

3596        if not hasattr(self, '_decisions') or decision_id not in self._decisions:3597            raise ValueError(f"Decision {decision_id} not found")35983599        try:3600            # Use graph traversal to find causal relationships3601            causal_chain = []3602            chain_limit = float("inf") if max_chains is None else max_chains3603            truncated = False36043605            # Reverse index of explicit causal edges, built once per call so the3606            # traversal does not rescan the edge list at every visited node.3607            # Edges may reference decision nodes that were never recorded through3608            # record_decision() (e.g. a graph restored via from_dict), so only3609            # causes with a known decision record are kept.3610            incoming_causal_edges = defaultdict(list)3611            # The index is keyed by the raw edge_type string ("causes" and3612            # "CAUSED" are separate keys), so filter by normalized type.3613            for edge_type, edges in self.edge_type_index.items():3614                if edge_type.upper() not in _CAUSAL_TRAVERSAL_TYPES:3615                    continue3616                for edge in edges:3617                    if edge.source_id in self._decisions:3618                        incoming_causal_edges[edge.target_id].append(edge)36193620            def record_chain(cause_path):3621                """Record one chain. Returns False once the cap is reached."""3622                nonlocal truncated3623                if len(causal_chain) >= chain_limit:3624                    truncated = True3625                    return False3626                causal_chain.append(3627                    self._build_causal_chain_report(list(reversed(cause_path)))3628                )3629                return True36303631            def trace_recursive(current_id, depth, path, path_ids):3632                # Cycle detection is per-path rather than global: a decision reached3633                # through one branch must stay traversable through another, otherwise3634                # branching graphs silently lose valid chains. max_depth bounds the3635                # traversal.3636                if truncated or depth >= max_depth or current_id in path_ids:3637                    return36383639                path_ids = path_ids | {current_id}3640                current_decision = self._decisions[current_id]36413642                # Explicit causal relationships recorded via add_causal_relationship()3643                # take precedence - they are the ground truth the caller recorded.3644                # Every edge is traced, so parallel relationships between the same3645                # pair of decisions are all reported rather than overwriting.3646                explicit_causes = incoming_causal_edges.get(current_id, [])3647                explicit_cause_ids = {edge.source_id for edge in explicit_causes}36483649                for edge in explicit_causes:3650                    cause_id = edge.source_id3651                    cause_dec = self._decisions[cause_id]3652                    weight = getattr(edge, "weight", None)3653                    # A stored weight of 0.0 is meaningful and must not be coerced3654                    # to the 1.0 default.3655                    edge_weight = 1.0 if weight is None else float(weight)3656                    hop = {3657                        "from": cause_id,3658                        "from_scenario": cause_dec.get("scenario", ""),3659                        "to": current_id,3660                        "to_scenario": current_decision.get("scenario", ""),3661                        "type": edge.edge_type,3662                        "edge_weight": edge_weight,3663                    }3664                    cause_path = path + [hop]3665                    if not record_chain(cause_path):3666                        return3667                    trace_recursive(cause_id, depth + 1, cause_path, path_ids)3668                    if truncated:3669                        return

3671                # Find potential causes (decisions that influenced this one) via3672                # shared entities/timestamps - additive heuristic, skipping anything3673                # already covered by an explicit relationship above.3674                potential_causes = []3675                for entity in current_decision["entities"]:3676                    for other_decision_id in self._entity_index.get(entity, set()):3677                        if other_decision_id != current_id and other_decision_id not in explicit_cause_ids:3678                            other_decision = self._decisions[other_decision_id]3679                            if other_decision["timestamp"] < current_decision["timestamp"]:3680                                potential_causes.append(other_decision_id)36813682                for cause_id in potential_causes:3683                    cause_dec = self._decisions.get(cause_id, {})3684                    edge_weight = float(cause_dec.get("confidence", 1.0))3685                    hop = {3686                        "from": cause_id,3687                        "from_scenario": cause_dec.get("scenario", ""),3688                        "to": current_id,3689                        "to_scenario": current_decision.get("scenario", ""),3690                        "type": "influences",3691                        "edge_weight": edge_weight,3692                    }3693                    cause_path = path + [hop]3694                    if not record_chain(cause_path):3695                        return

91    def get_causal_chain(92        self,93        decision_id: str,94        direction: str = "upstream",95        max_depth: int = 1096    ) -> List[Decision]:97        """98        Trace decision causality in specified direction.99        100        Args:101            decision_id: Starting decision ID102            direction: "upstream" (what caused this) or "downstream" (what this caused)103            max_depth: Maximum traversal depth104            105        Returns:106            List of decisions in causal chain107        """108        try:109            if hasattr(self.graph_store, "get_causal_chain") and not hasattr(self.graph_store, "execute_query"):110                return self.graph_store.get_causal_chain(111                    decision_id=decision_id,112                    direction=direction,113                    max_depth=max_depth114                )115116            if not (1 <= max_depth <= 100):117                raise ValueError("max_depth must be between 1 and 20")118119            if direction not in ["upstream", "downstream"]:120                raise ValueError("Direction must be 'upstream' or 'downstream'")121            122            # Define relationship direction based on traversal direction123            if direction == "upstream":124                rel_pattern = "<-[:CAUSED|:INFLUENCED|:PRECEDENT_FOR]-"125            else:126                rel_pattern = "-[:CAUSED|:INFLUENCED|:PRECEDENT_FOR]->"127            128            query = f"""129            MATCH (start:Decision {{decision_id: $decision_id}})130            MATCH path = (start){rel_pattern}{{1,{max_depth}}}(end:Decision)131            RETURN DISTINCT end, length(path) as distance132            ORDER BY distance, end.timestamp133            """134            135            results = self.graph_store.execute_query(query, {136                "decision_id": decision_id137            })138            results = self._extract_records(results)139            140            decisions = []141            for record in results:142                decision_data = record.get("end") if isinstance(record, dict) else None143                if not isinstance(decision_data, dict):144                    decision_data = record if isinstance(record, dict) else {}145                decision = self._dict_to_decision(decision_data)146                decision.metadata["causal_distance"] = record.get("distance", 0)147                decisions.append(decision)148            149            self.logger.info(f"Found {len(decisions)} decisions in {direction} causal chain")150            return decisions151            152        except Exception as e:153            self.logger.error(f"Failed to get causal chain: {e}")154            raise

72class CausalChainAnalyzer:73    """74    Analyzes causal chains between decisions.75    76    This class provides methods for tracing decision causality, finding77    decisions that influenced others, and analyzing precedent relationships78    using graph traversal.79    """80    81    def __init__(self, graph_store: Any):82        """83        Initialize CausalChainAnalyzer.84        85        Args:86            graph_store: Graph database instance for traversal87        """88        self.graph_store = graph_store89        self.logger = get_logger(__name__)

5The context module gives your AI agents the ability to **remember**, **learn**, and **make smarter decisions** by organizing information in a way that's both powerful and easy to use.67Think of it as giving your agent a brain that can:8- **Remember conversations** (like human memory)9- **Learn from past decisions** (become smarter over time)10- **Find relevant information** quickly (when it matters most)11- **Understand relationships** between concepts12- **Make consistent decisions** based on experience

174@dataclass175class Policy:176    """Policy with versioning and change tracking."""177    178    policy_id: str179    name: str180    description: str181    rules: Dict[str, Any]182    category: str183    version: str184    created_at: datetime185    updated_at: datetime186    metadata: Dict[str, Any] = field(default_factory=dict)187    188    def __post_init__(self, auto_generate_id: bool = True):189        """Validate policy data."""190        if auto_generate_id and not self.policy_id:  # Handle both None and empty string191            self.policy_id = str(uuid.uuid4())192        elif not self.policy_id and not auto_generate_id:193            raise ValueError("policy_id is required when auto_generate_id=False")194    195    def to_dict(self) -> Dict[str, Any]:196        """Convert policy to dictionary."""197        return {198            "policy_id": self.policy_id,199            "name": self.name,200            "description": self.description,201            "rules": self.rules,202            "category": self.category,203            "version": self.version,204            "created_at": self.created_at.isoformat(),205            "updated_at": self.updated_at.isoformat(),206            "metadata": self.metadata207        }208    209    @classmethod210    def from_dict(cls, data: Dict[str, Any]) -> "Policy":211        """Create policy from dictionary."""212        for field in ["created_at", "updated_at"]:213            if isinstance(data.get(field), str):214                data[field] = datetime.fromisoformat(data[field])215        return cls(**data)

79class PolicyEngine:80    """81    Policy engine with versioning and change tracking.82    83    This class manages policies, tracks versions, checks compliance,84    records exceptions, and analyzes policy impact.85    """86    87    def __init__(self, graph_store: Any):88        """89        Initialize PolicyEngine.90        91        Args:92            graph_store: Graph database instance for storing policies93        """94        self.graph_store = graph_store95        self.logger = get_logger(__name__)96        self._supports_cypher = hasattr(graph_store, "execute_query")

167    def update_policy(168        self,169        policy_id: str,170        rules: Dict[str, Any],171        change_reason: str,172        new_version: Optional[str] = None173    ) -> str:174        """175        Update policy and create new version.176        177        Args:178            policy_id: Policy ID to update179            rules: New policy rules180            change_reason: Reason for the change181            new_version: Optional new version (auto-generated if not provided)182            183        Returns:184            New version string185        """186        try:187            # Get current policy188            current_policy = self.get_policy(policy_id)189            if not current_policy:190                raise ValueError(f"Policy {policy_id} not found")191            192            # Generate new version if not provided193            if not new_version:194                new_version = self._generate_next_version(current_policy.version)195            196            # Create new policy version197            updated_policy = Policy(198                policy_id=policy_id,199                name=current_policy.name,200                description=current_policy.description,201                rules=rules,202                category=current_policy.category,203                version=new_version,204                created_at=current_policy.created_at,205                updated_at=datetime.now(),206                metadata={207                    **current_policy.metadata,208                    "change_reason": change_reason,209                    "previous_version": current_policy.version210                }211            )212            213            # Store new version214            self.add_policy(updated_policy)215            216            if self._supports_cypher:217                query = """218                MATCH (old:Policy {policy_id: $policy_id, version: $old_version})219                MATCH (new:Policy {policy_id: $policy_id, version: $new_version})220                MERGE (old)-[:VERSION_OF]->(new)221                """222                self.graph_store.execute_query(query, {223                    "policy_id": policy_id,224                    "old_version": current_policy.version,225                    "new_version": new_version226                })227            else:228                if hasattr(self.graph_store, "add_edge"):229                    self.graph_store.add_edge(230                        f"{policy_id}:{current_policy.version}",231                        f"{policy_id}:{new_version}",232                        edge_type="VERSION_OF",233                        changed_at=datetime.now().isoformat(),234                        change_reason=change_reason235                    )236            237            self.logger.info(f"Updated policy {policy_id} to version {new_version}")238            return policy_id239            240        except Exception as e:241            self.logger.exception("Failed to update policy")242            raise

896    def _evaluate_compliance(self, decision: Decision, rules: Dict[str, Any]) -> bool:897        """Evaluate decision compliance against policy rules."""898        metadata = decision.metadata or {}899900        for rule_key, rule_value in rules.items():901            # min_X → metadata["X"] >= rule_value902            if rule_key.startswith("min_"):903                field = rule_key[4:]904                field_value = self._get_metadata_field(metadata, decision, field)905                if field_value is None:906                    return False907                if field_value < rule_value:908                    return False909            # max_X → metadata["X"] <= rule_value910            elif rule_key.startswith("max_"):911                field = rule_key[4:]912                field_value = self._get_metadata_field(metadata, decision, field)913                if field_value is None:914                    return False915                # For strings use lexicographic comparison916                if field_value > rule_value:917                    return False918            # required_X → metadata["X"] contains all items in rule_value (list) or equals (str)919            elif rule_key.startswith("required_"):920                field = rule_key[9:]921                field_value = self._get_metadata_field(metadata, decision, field)922                if field_value is None:923                    return False924                if isinstance(rule_value, list):925                    if not all(item in field_value for item in rule_value):926                        return False927                elif field_value != rule_value:928                    return False929            # Direct checks from _check_compliance_with_rules930            elif rule_key in {"min_confidence", "allowed_outcomes", "required_categories"}:931                decision_data = {"confidence": decision.confidence, "outcome": decision.outcome,932                                 "category": decision.category}933                if not self._check_compliance_with_rules(decision_data, {rule_key: rule_value}):934                    return False935            else:936                # Unknown rule key — check if field is present in metadata937                if rule_key not in metadata:938                    return False939        return True

3726    def enforce_decision_policy(3727        self,3728        decision_data: Dict[str, Any],3729        policy_rules: Optional[Dict[str, Any]] = None3730    ) -> Dict[str, Any]:3731        """3732        Enforce policies on decision data.3733        3734        Args:3735            decision_data: Decision data to check3736            policy_rules: Policy rules to enforce3737            3738        Returns:3739            Policy enforcement results3740        """3741        # Simple policy enforcement implementation3742        violations = []3743        warnings = []3744        3745        # Default policy rules3746        default_rules = {3747            "min_confidence": 0.7,3748            "required_outcomes": ["approved", "rejected", "flagged"],3749            "required_metadata": ["decision_maker"],3750            "max_reasoning_length": 10003751        }3752        3753        rules = policy_rules or default_rules3754        3755        # Check confidence3756        if decision_data.get("confidence", 0) < rules.get("min_confidence", 0.7):3757            violations.append(f"Confidence too low: {decision_data.get('confidence', 0)}")3758        3759        # Check outcome3760        if decision_data.get("outcome") not in rules.get("required_outcomes", []):3761            violations.append(f"Invalid outcome: {decision_data.get('outcome')}")3762        3763        # Check required metadata3764        for required_field in rules.get("required_metadata", []):3765            if not decision_data.get(required_field):3766                violations.append(f"Missing required field: {required_field}")3767        3768        # Check reasoning length3769        reasoning = decision_data.get("reasoning", "")3770        if len(reasoning) > rules.get("max_reasoning_length", 1000):3771            warnings.append(f"Reasoning too long: {len(reasoning)} characters")3772        3773        return {3774            "compliant": len(violations) == 0,3775            "violations": violations,3776            "warnings": warnings,3777            "policy_rules": rules3778        }

429    for event in events:430        event_type = event.get("event_type", "TRACE_EVENT")431        payload = event.get("payload", {})432        payload_json = json.dumps(payload, sort_keys=True, default=str)433        event_timestamp = datetime.now().isoformat()434        trace_id = f"{decision_id}:{next_index}"435        hash_input = (436            f"{decision_id}|{next_index}|{event_type}|{event_timestamp}|{payload_json}|{previous_hash}"437        )438        event_hash = hashlib.sha256(hash_input.encode("utf-8")).hexdigest()439440        graph_store.execute_query(441            """442            MATCH (d:Decision {decision_id: $decision_id})443            CREATE (t:DecisionTraceEvent {444                trace_id: $trace_id,445                decision_id: $decision_id,446                event_index: $event_index,447                event_type: $event_type,448                event_timestamp: $event_timestamp,449                event_payload: $event_payload,450                previous_hash: $previous_hash,451                event_hash: $event_hash452            })453            MERGE (d)-[:HAS_TRACE_EVENT]->(t)454            """,455            {456                "decision_id": decision_id,457                "trace_id": trace_id,458                "event_index": next_index,459                "event_type": event_type,460                "event_timestamp": event_timestamp,461                "event_payload": payload_json,462                "previous_hash": previous_hash,463                "event_hash": event_hash,464            },465        )

339            "DecisionTraceEvent": {340                "properties": [341                    "trace_id", "decision_id", "event_index", "event_type",342                    "event_timestamp", "event_payload", "previous_hash", "event_hash"343                ],344                "constraints": ["decision_trace_id_unique"],345                "indexes": [346                    "decision_trace_id_index",347                    "decision_trace_event_index",348                    "decision_trace_type_index",

56    async def search(57        self,58        query: str,59        conversation_history: ConversationHistory | None = None,60        **kwargs,61    ) -> SearchResult:62        """Build local search context that fits a single context window and generate answer for the user query."""63        start_time = time.time()64        search_prompt = ""65        llm_calls, prompt_tokens, output_tokens = {}, {}, {}66        context_result = self.context_builder.build_context(67            query=query,68            conversation_history=conversation_history,69            **kwargs,70            **self.context_builder_params,71        )72        llm_calls["build_context"] = context_result.llm_calls73        prompt_tokens["build_context"] = context_result.prompt_tokens74        output_tokens["build_context"] = context_result.output_tokens

92            messages_builder = (93                CompletionMessagesBuilder()94                .add_system_message(search_prompt)95                .add_user_message(query)96            )9798            full_response = ""99100            response: AsyncIterator[101                LLMCompletionChunk102            ] = await self.model.completion_async(103                messages=messages_builder.build(),104                stream=True,105                **self.model_params,106            )  # type: ignore107108            async for chunk in response:109                response_text = chunk.choices[0].delta.content or ""110                full_response += response_text111                for callback in self.callbacks:112                    callback.on_llm_new_token(response_text)113114            llm_calls["response"] = 1115            prompt_tokens["response"] = len(self.tokenizer.encode(search_prompt))116            output_tokens["response"] = len(self.tokenizer.encode(full_response))117118            for callback in self.callbacks:119                callback.on_context(context_result.context_records)120121            return SearchResult(

30def build_entity_context(31    selected_entities: list[Entity],32    tokenizer: Tokenizer | None = None,33    max_context_tokens: int = 8000,34    include_entity_rank: bool = True,35    rank_description: str = "number of relationships",36    column_delimiter: str = "|",37    context_name="Entities",38) -> tuple[str, pd.DataFrame]:39    """Prepare entity data table as context data for system prompt."""40    tokenizer = tokenizer or get_tokenizer()4142    if len(selected_entities) == 0:43        return "", pd.DataFrame()4445    # add headers46    current_context_text = f"-----{context_name}-----" + "\n"47    header = ["id", "entity", "description"]48    if include_entity_rank:49        header.append(rank_description)50    attribute_cols = (51        list(selected_entities[0].attributes.keys())52        if selected_entities[0].attributes53        else []54    )55    header.extend(attribute_cols)56    current_context_text += column_delimiter.join(header) + "\n"57    current_tokens = tokenizer.num_tokens(current_context_text)