Janady13/AI_Brain / services/knowledge_pipeline.py

1"""Lightweight knowledge orchestration demo covering RAG and KAG flows. 2 3The repository’s marketing ANI exposes rich process payloads for retrieval- 4augmented generation (RAG) and knowledge-augmented generation (KAG). The 5production services already synthesise these insights, but it can be handy to 6have a standalone executable module that demonstrates the core mechanics with 7minimal dependencies. This module keeps the implementation intentionally 8simple so it can run in constrained environments while still mimicking the 9steps outlined in the documentation: 10 11* **Indexing** – documents are split into manageable chunks and converted into 12 term-frequency vectors ready for fast cosine similarity searches. 13* **Retrieval** – queries are embedded with the same lightweight tokenizer and 14 matched against the indexed chunks to surface relevant context, giving us a 15 local RAG demo without external services. 16* **Knowledge graph reasoning** – structured facts are captured in a small 17 in-memory graph so we can traverse relationships and ground answers with 18 explicit entity links, mirroring a KAG flow. 19* **Orchestration** – the :class:`KnowledgeWorkflowOrchestrator` blends both 20 approaches, returning combined context and reasoning trails that resemble the 21 payload emitted by the marketing ANI’s processes endpoint. 22 23The goal is to make it easy for contributors (and automated smoke tests) to 24verify that the conceptual pipelines run end-to-end. The code favours 25deterministic behaviour over sophistication so that unit tests can assert on 26exact outputs without relying on heavyweight ML frameworks. 27""" 28 29from __future__ import annotations 30 31import argparse 32import json 33import math 34import os 35import re 36from collections import Counter, defaultdict 37from dataclasses import dataclass, asdict 38from typing import Dict, Iterator, List, Mapping, MutableMapping, Sequence 39 40 41_TOKEN_PATTERN = re.compile(r"[A-Za-z0-9]+") 42 43 44def _tokenize(text: str) -> List[str]: 45 """Tokenise ``text`` into lowercase word tokens. 46 47 The helper intentionally keeps the implementation minimal so the demo does 48 not rely on heavy NLP libraries. Tokens are restricted to alphanumeric 49 characters which keeps outputs predictable for unit tests. 50 """ 51 52 return [token.lower() for token in _TOKEN_PATTERN.findall(text)] 53 54 55def _normalise(counter: Counter[str]) -> Dict[str, float]: 56 """Convert a raw term frequency ``counter`` into L2-normalised weights.""" 57 58 magnitude = math.sqrt(sum(value * value for value in counter.values())) 59 if magnitude == 0.0: 60 return {} 61 return {key: value / magnitude for key, value in counter.items()} 62 63 64def _cosine_similarity(lhs: Mapping[str, float], rhs: Mapping[str, float]) -> float: 65 """Compute cosine similarity between two sparse vectors.""" 66 67 if not lhs or not rhs: 68 return 0.0 69 if len(lhs) > len(rhs): 70 lhs, rhs = rhs, lhs 71 return sum(weight * rhs.get(term, 0.0) for term, weight in lhs.items()) 72 73 74def _chunk_text(text: str, *, max_tokens: int) -> Iterator[str]: 75 """Yield whitespace-trimmed chunks that do not exceed ``max_tokens`` tokens.""" 76 77 sentences = re.split(r"(?<=[.!?])\s+", text.strip()) if text.strip() else [] 78 buffer: List[str] = [] 79 token_count = 0 80 for sentence in sentences: 81 sentence_tokens = _tokenize(sentence) 82 if not sentence_tokens: 83 continue 84 if token_count + len(sentence_tokens) > max_tokens and buffer: 85 yield " ".join(buffer).strip() 86 buffer = [sentence] 87 token_count = len(sentence_tokens) 88 else: 89 buffer.append(sentence) 90 token_count += len(sentence_tokens) 91 if buffer: 92 yield " ".join(buffer).strip() 93 94 95@dataclass(frozen=True) 96class DocumentChunk: 97 """Indexed document chunk used by the :class:`RAGPipeline`.""" 98 99 doc_id: str 100 chunk_index: int 101 text: str 102 vector: Mapping[str, float] 103 104 105class RAGPipeline: 106 """Very small RAG pipeline used for local demonstrations and tests.""" 107 108 def __init__(self, *, chunk_token_limit: int = 120) -> None: 109 self._chunk_token_limit = chunk_token_limit 110 self._chunks: List[DocumentChunk] = [] 111 self._cache: MutableMapping[str, List[DocumentChunk]] = {} 112 113 @property 114 def chunks(self) -> Sequence[DocumentChunk]: 115 return tuple(self._chunks) 116 117 def index(self, documents: Mapping[str, str]) -> None: 118 """Index ``documents`` ready for retrieval queries.""" 119 120 self._chunks.clear() 121 self._cache.clear() 122 for doc_id, text in documents.items(): 123 for chunk_index, chunk_text in enumerate( 124 _chunk_text(text, max_tokens=self._chunk_token_limit) 125 ): 126 tokens = Counter(_tokenize(chunk_text)) 127 if not tokens: 128 continue 129 self._chunks.append( 130 DocumentChunk( 131 doc_id=doc_id, 132 chunk_index=chunk_index, 133 text=chunk_text, 134 vector=_normalise(tokens), 135 ) 136 ) 137 138 def search(self, query: str, *, limit: int = 3) -> List[DocumentChunk]: 139 """Return up to ``limit`` document chunks relevant to ``query``.""" 140 141 if query in self._cache: 142 return list(self._cache[query][:limit]) 143 144 query_tokens = Counter(_tokenize(query)) 145 query_vector = _normalise(query_tokens) 146 scored: List[tuple[float, DocumentChunk]] = [] 147 for chunk in self._chunks: 148 score = _cosine_similarity(chunk.vector, query_vector) 149 if score <= 0.0: 150 continue 151 scored.append((score, chunk)) 152 scored.sort(key=lambda item: item[0], reverse=True) 153 top_chunks = [chunk for _, chunk in scored[:limit]] 154 self._cache[query] = top_chunks 155 return list(top_chunks) 156 157 158@dataclass(frozen=True) 159class KnowledgeEdge: 160 """Represents a directed edge in the demo knowledge graph.""" 161 162 source: str 163 relation: str 164 target: str 165 evidence: str | None = None 166 167 168class KnowledgeGraph: 169 """Extremely small in-memory knowledge graph with simple traversal.""" 170 171 def __init__(self) -> None: 172 self._graph: MutableMapping[str, List[KnowledgeEdge]] = defaultdict(list) 173 174 def add_fact( 175 self, 176 source: str, 177 relation: str, 178 target: str, 179 *, 180 evidence: str | None = None, 181 ) -> None: 182 edge = KnowledgeEdge(source=source, relation=relation, target=target, evidence=evidence) 183 self._graph[source].append(edge) 184 185 def related(self, entity: str, relation: str | None = None) -> List[KnowledgeEdge]: 186 edges = list(self._graph.get(entity, ())) 187 if relation is None: 188 return edges 189 return [edge for edge in edges if edge.relation == relation] 190 191 def reason(self, query: str) -> List[KnowledgeEdge]: 192 """Retrieve graph edges that align with ``query`` tokens.""" 193 194 tokens = set(_tokenize(query)) 195 matched: List[KnowledgeEdge] = [] 196 for edges in self._graph.values(): 197 for edge in edges: 198 edge_tokens = set(_tokenize(edge.source)) | set(_tokenize(edge.relation)) | set( 199 _tokenize(edge.target) 200 ) 201 if edge_tokens & tokens: 202 matched.append(edge) 203 return matched 204 205 206@dataclass(frozen=True) 207class KnowledgeWorkflowAnswer: 208 """Combined answer produced by orchestrating RAG and KAG flows.""" 209 210 rag_context: List[str] 211 knowledge_edges: List[KnowledgeEdge] 212 213 def to_dict(self) -> Dict[str, object]: 214 return { 215 "rag_context": list(self.rag_context), 216 "knowledge_edges": [asdict(edge) for edge in self.knowledge_edges], 217 } 218 219 220class KnowledgeWorkflowOrchestrator: 221 """Fuse RAG and KAG results into a single structured response.""" 222 223 def __init__(self, rag: RAGPipeline, graph: KnowledgeGraph) -> None: 224 self._rag = rag 225 self._graph = graph 226 227 def answer(self, query: str, *, limit: int = 3) -> KnowledgeWorkflowAnswer: 228 rag_chunks = self._rag.search(query, limit=limit) 229 edges = self._graph.reason(query) 230 return KnowledgeWorkflowAnswer( 231 rag_context=[chunk.text for chunk in rag_chunks], 232 knowledge_edges=edges, 233 ) 234 235 236def _load_default_documents() -> Dict[str, str]: 237 """Load demonstration documents that complement the marketing ANI datasets.""" 238 239 base_dir = os.path.join(os.path.dirname(__file__), "..", "daemon", "datasets") 240 dataset_path = os.path.abspath(os.path.join(base_dir, "knowledge_workflows.json")) 241 documents: Dict[str, str] = {} 242 try: 243 with open(dataset_path, "r", encoding="utf-8") as handle: 244 payload = json.load(handle) 245 except FileNotFoundError: 246 return documents 247 248 for plan in payload: 249 key = f"{plan['category']}_{plan['name']}" 250 description = plan.get("description") or "" 251 indexing_steps = plan.get("indexing_steps") or [] 252 retrieval_steps = plan.get("retrieval_steps") or [] 253 governance = plan.get("governance") or [] 254 evaluation = plan.get("evaluation") or [] 255 content_parts = [ 256 description, 257 " ".join(indexing_steps), 258 " ".join(retrieval_steps), 259 " ".join(governance), 260 " ".join(evaluation), 261 ] 262 documents[key] = " ".join(part for part in content_parts if part) 263 return documents 264 265 266def _build_demo_graph() -> KnowledgeGraph: 267 graph = KnowledgeGraph() 268 graph.add_fact( 269 "retrieval_pipeline", 270 "uses", 271 "vector_database", 272 evidence="Vector search keeps embeddings queryable for RAG.", 273 ) 274 graph.add_fact( 275 "knowledge_graph", 276 "aligns_with", 277 "prompt_diagnostics", 278 evidence="Structured triples link prompts to governance policies.", 279 ) 280 graph.add_fact( 281 "agent_loop", 282 "includes", 283 "reflection", 284 evidence="Post-action analysis keeps the ANI adaptive.", 285 ) 286 graph.add_fact( 287 "moe_router", 288 "activates", 289 "stripe", 290 evidence="Stripe expert handles monetisation strategies.", 291 ) 292 graph.add_fact( 293 "moe_router", 294 "activates", 295 "google_ads", 296 evidence="Google Ads expert supplies paid acquisition intelligence.", 297 ) 298 return graph 299 300 301def run_demo(query: str, *, limit: int = 3) -> KnowledgeWorkflowAnswer: 302 """Execute the lightweight RAG+KAG demo for ``query``.""" 303 304 rag = RAGPipeline() 305 rag.index(_load_default_documents()) 306 graph = _build_demo_graph() 307 orchestrator = KnowledgeWorkflowOrchestrator(rag=rag, graph=graph) 308 return orchestrator.answer(query, limit=limit) 309 310 311def _parse_args() -> argparse.Namespace: 312 parser = argparse.ArgumentParser(description="Run the lightweight RAG/KAG demo") 313 parser.add_argument( 314 "query", 315 nargs="?", 316 default="How do the RAG and KAG workflows keep monetisation grounded?", 317 help="Natural language query to execute", 318 ) 319 parser.add_argument("--limit", type=int, default=3, help="Maximum RAG chunks to return") 320 return parser.parse_args() 321 322 323def main() -> None: 324 args = _parse_args() 325 answer = run_demo(args.query, limit=args.limit) 326 print(json.dumps(answer.to_dict(), indent=2)) 327 328 329if __name__ == "__main__": # pragma: no cover - exercised via smoke usage 330 main() 331