RAG-based knowledge base for storing and retrieving guidance. Storage format: JSONL with {subject: str, guidance: str, embedding: List[float]}
| 13 | |
| 14 | |
| 15 | class KnowledgeBase: |
| 16 | """ |
| 17 | RAG-based knowledge base for storing and retrieving guidance. |
| 18 | |
| 19 | Storage format: JSONL with {subject: str, guidance: str, embedding: List[float]} |
| 20 | """ |
| 21 | |
| 22 | def __init__( |
| 23 | self, |
| 24 | storage_path: str, |
| 25 | llm_client: LLMClient, |
| 26 | max_guidance_length: int = 1024, |
| 27 | guidance_merge_prompt_template: Optional[str] = None, |
| 28 | ): |
| 29 | """ |
| 30 | Initialize the knowledge base. |
| 31 | |
| 32 | Args: |
| 33 | storage_path: Path to JSONL storage file |
| 34 | llm_client: LLM client for embeddings and merging |
| 35 | max_guidance_length: Maximum character length for guidance |
| 36 | guidance_merge_prompt_template: Custom template for merging guidance |
| 37 | """ |
| 38 | self.storage_path = storage_path |
| 39 | self.llm_client = llm_client |
| 40 | self.max_guidance_length = max_guidance_length |
| 41 | self.guidance_merge_prompt_template = guidance_merge_prompt_template or self._get_default_merge_prompt_template() |
| 42 | |
| 43 | # Load existing knowledge base |
| 44 | self.entries = self._load_entries() |
| 45 | |
| 46 | def _load_entries(self) -> List[Dict[str, Any]]: |
| 47 | """Load entries from storage file.""" |
| 48 | if self.storage_path and os.path.exists(self.storage_path): |
| 49 | return load_jsonl(self.storage_path) |
| 50 | return [] |
| 51 | |
| 52 | def _get_default_merge_prompt_template(self) -> str: |
| 53 | """Get the default prompt template for merging guidance.""" |
| 54 | return ( |
| 55 | "You are synthesizing guidance for the subject: {subject}\n\n" |
| 56 | "Existing guidance from related subjects in the knowledge base:\n" |
| 57 | "{existing_guidance}\n\n" |
| 58 | "New guidance to incorporate:\n{new_guidance}\n\n" |
| 59 | "Please merge these guidance points into a single, coherent guidance text for '{subject}'.\n" |
| 60 | "Consider insights from related subjects and adapt them to the current context.\n" |
| 61 | "The output should be concise, clear, and no longer than {max_length} characters.\n" |
| 62 | "Focus on the most important and actionable advice.\n\n" |
| 63 | "Merged guidance:" |
| 64 | ) |
| 65 | |
| 66 | def _save_entries(self) -> None: |
| 67 | """Save entries to storage file.""" |
| 68 | save_jsonl(self.entries, self.storage_path, append=False) |
| 69 | |
| 70 | def retrieve( |
| 71 | self, |
| 72 | query: str, |