myagent-ai 1.0.0

package/knowledge/__init__.py

@@ -0,0 +1,4 @@
+"""
+knowledge/__init__.py - 知识库模块入口
+"""
+from knowledge.rag import KnowledgeRAG

package/knowledge/rag.py

@@ -0,0 +1,444 @@
+"""
+knowledge/rag.py - 轻量级 RAG 知识检索模块
+==========================================
+基于 TF-IDF 的文本检索系统，支持:
+  - 读取知识库目录中的文本文件
+  - 文本分块 (chunk_size, overlap)
+  - TF-IDF 向量化与余弦相似度搜索
+  - 支持组织级别和 Agent 级别的知识库
+
+支持文件类型: .md, .txt, .json, .csv, .py, .js, .html
+"""
+from __future__ import annotations
+
+import json
+import math
+import os
+import re
+from collections import Counter
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from core.logger import get_logger
+from core.utils import timestamp
+
+logger = get_logger("myagent.knowledge.rag")
+
+# 支持的文件扩展名
+SUPPORTED_EXTENSIONS = {".md", ".txt", ".json", ".csv", ".py", ".js", ".html"}
+
+
+@dataclass
+class KnowledgeChunk:
+    """知识文本块"""
+    chunk_id: str = ""
+    file_path: str = ""         # 相对路径
+    file_name: str = ""
+    content: str = ""
+    chunk_index: int = 0        # 在原文件中的块索引
+    score: float = 0.0          # 搜索得分
+
+    def to_dict(self) -> dict:
+        return {
+            "chunk_id": self.chunk_id,
+            "file_path": self.file_path,
+            "file_name": self.file_name,
+            "content": self.content[:2000],  # 限制返回长度
+            "chunk_index": self.chunk_index,
+            "score": round(self.score, 4),
+        }
+
+
+class KnowledgeRAG:
+    """
+    轻量级 RAG 知识检索引擎。
+
+    使用示例:
+        rag = KnowledgeRAG(kb_dir="~/.myagent/data/organization/knowledge")
+        rag.build_index()
+        results = rag.search("公司报销流程", top_k=5)
+    """
+
+    def __init__(
+        self,
+        kb_dir: str | Path = "",
+        chunk_size: int = 500,
+        overlap: int = 50,
+    ):
+        """
+        Args:
+            kb_dir: 知识库目录路径
+            chunk_size: 文本分块大小（字符数）
+            overlap: 分块重叠大小（字符数）
+        """
+        self.kb_dir = Path(kb_dir) if kb_dir else Path()
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+        # 索引数据: {chunk_id: KnowledgeChunk}
+        self._chunks: Dict[str, KnowledgeChunk] = {}
+        # 原始文档缓存: {相对路径: 完整内容}
+        self._documents: Dict[str, str] = {}
+
+    # ==========================================================================
+    # TF-IDF 分词（复用 memory/manager.py 的中文分词策略）
+    # ==========================================================================
+
+    @staticmethod
+    def _tokenize(text: str) -> List[str]:
+        """
+        中文分词（简单实现：单字+双字组合）+ 英文词提取。
+        """
+        if not text:
+            return []
+
+        tokens: List[str] = []
+        # 提取中文连续片段
+        chinese_segments = re.findall(r'[\u4e00-\u9fff]+', text)
+        for seg in chinese_segments:
+            tokens.extend(list(seg))
+            for i in range(len(seg) - 1):
+                tokens.append(seg[i:i + 2])
+
+        # 提取英文/数字单词
+        english_words = re.findall(r'[a-zA-Z0-9]+', text.lower())
+        tokens.extend(english_words)
+
+        return tokens
+
+    @staticmethod
+    def _compute_tfidf(
+        query: str,
+        documents: List[Tuple[str, str]],
+    ) -> Dict[str, float]:
+        """
+        计算 TF-IDF 相似度得分。
+
+        Args:
+            query: 查询文本
+            documents: 文档列表 [(doc_id, text), ...]
+
+        Returns:
+            {doc_id: tfidf_score} 按得分降序排列
+        """
+        if not documents or not query:
+            return {}
+
+        _tokenize = KnowledgeRAG._tokenize
+
+        query_tokens = _tokenize(query)
+        if not query_tokens:
+            return {}
+
+        n_docs = len(documents)
+
+        # 计算每个文档的 TF
+        doc_tfs: Dict[str, Counter] = {}
+        for doc_id, text in documents:
+            doc_tfs[doc_id] = Counter(_tokenize(text))
+
+        # 计算 IDF
+        doc_freq: Counter = Counter()
+        for doc_id, tf in doc_tfs.items():
+            for token in set(tf.keys()):
+                doc_freq[token] += 1
+
+        idf: Dict[str, float] = {}
+        for token, df in doc_freq.items():
+            idf[token] = math.log((n_docs + 1) / (1 + df)) + 1
+
+        # 查询的 TF
+        query_tf = Counter(query_tokens)
+
+        # 计算查询向量的模
+        query_norm = math.sqrt(
+            sum((query_tf.get(t, 0) * idf.get(t, 0)) ** 2 for t in query_tf)
+        )
+        if query_norm == 0:
+            return {}
+
+        # 余弦相似度
+        scores: Dict[str, float] = {}
+        for doc_id in doc_tfs:
+            doc_tf = doc_tfs[doc_id]
+            dot_product = 0.0
+            for token in query_tf:
+                if token in doc_tf:
+                    dot_product += (query_tf[token] * idf.get(token, 0)) * \
+                                   (doc_tf[token] * idf.get(token, 0))
+
+            doc_norm = math.sqrt(
+                sum((doc_tf.get(t, 0) * idf.get(t, 0)) ** 2 for t in doc_tf)
+            ) if doc_tf else 0
+
+            if doc_norm > 0:
+                scores[doc_id] = dot_product / (query_norm * doc_norm)
+            else:
+                scores[doc_id] = 0.0
+
+        return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True))
+
+    # ==========================================================================
+    # 文本分块
+    # ==========================================================================
+
+    def _split_text(self, text: str) -> List[str]:
+        """
+        将文本按段落分块，保持语义完整性。
+
+        优先在段落边界（空行）处分割，不足 chunk_size 时按字符切分。
+        """
+        if not text or not text.strip():
+            return []
+
+        # 按空行分段
+        paragraphs = re.split(r'\n\s*\n', text)
+        paragraphs = [p.strip() for p in paragraphs if p.strip()]
+
+        if not paragraphs:
+            return []
+
+        chunks: List[str] = []
+        current_chunk = ""
+
+        for para in paragraphs:
+            if not current_chunk:
+                current_chunk = para
+            elif len(current_chunk) + len(para) + 2 <= self.chunk_size:
+                current_chunk += "\n\n" + para
+            else:
+                # 当前块已满，先保存
+                chunks.append(current_chunk)
+                # 如果段落本身超过 chunk_size，按行切分
+                if len(para) > self.chunk_size:
+                    # 先保存已有部分
+                    lines = para.split('\n')
+                    current_chunk = ""
+                    for line in lines:
+                        if not current_chunk:
+                            current_chunk = line
+                        elif len(current_chunk) + len(line) + 1 <= self.chunk_size:
+                            current_chunk += "\n" + line
+                        else:
+                            chunks.append(current_chunk)
+                            # 重叠: 保留最后 overlap 个字符
+                            if self.overlap > 0 and len(current_chunk) > self.overlap:
+                                current_chunk = current_chunk[-self.overlap:] + "\n" + line
+                            else:
+                                current_chunk = line
+                else:
+                    # 重叠: 保留前一块的末尾
+                    if self.overlap > 0 and len(current_chunk) > self.overlap:
+                        current_chunk = current_chunk[-self.overlap:] + "\n\n" + para
+                    else:
+                        current_chunk = para
+
+        if current_chunk:
+            chunks.append(current_chunk)
+
+        return chunks
+
+    # ==========================================================================
+    # 索引管理
+    # ==========================================================================
+
+    def _read_file_content(self, file_path: Path) -> str:
+        """读取文件内容，支持多种文本格式"""
+        try:
+            content = file_path.read_text(encoding="utf-8", errors="replace")
+            # JSON 文件提取纯文本
+            if file_path.suffix == ".json":
+                try:
+                    obj = json.loads(content)
+                    content = json.dumps(obj, ensure_ascii=False, indent=2)
+                except json.JSONDecodeError:
+                    pass
+            return content
+        except Exception as e:
+            logger.warning(f"读取文件失败 {file_path}: {e}")
+            return ""
+
+    def build_index(self) -> int:
+        """
+        扫描知识库目录，构建全文索引。
+
+        Returns:
+            索引的块数量
+        """
+        if not self.kb_dir.exists():
+            logger.warning(f"知识库目录不存在: {self.kb_dir}")
+            return 0
+
+        self._chunks.clear()
+        self._documents.clear()
+
+        file_count = 0
+        chunk_count = 0
+
+        # 递归扫描文件
+        for file_path in sorted(self.kb_dir.rglob("*")):
+            if not file_path.is_file():
+                continue
+            if file_path.suffix.lower() not in SUPPORTED_EXTENSIONS:
+                continue
+
+            content = self._read_file_content(file_path)
+            if not content.strip():
+                continue
+
+            rel_path = str(file_path.relative_to(self.kb_dir))
+            self._documents[rel_path] = content
+
+            # 分块
+            text_chunks = self._split_text(content)
+            for idx, chunk_text in enumerate(text_chunks):
+                chunk_id = f"{rel_path}#chunk_{idx}"
+                kc = KnowledgeChunk(
+                    chunk_id=chunk_id,
+                    file_path=rel_path,
+                    file_name=file_path.name,
+                    content=chunk_text,
+                    chunk_index=idx,
+                )
+                self._chunks[chunk_id] = kc
+                chunk_count += 1
+
+            file_count += 1
+
+        logger.info(
+            f"知识库索引构建完成: {file_count} 个文件, {chunk_count} 个文本块 "
+            f"(dir={self.kb_dir})"
+        )
+        return chunk_count
+
+    def add_document(self, file_path: str, content: str = "") -> bool:
+        """
+        添加单个文档到索引。
+
+        Args:
+            file_path: 文件相对路径（相对于知识库根目录）
+            content: 文件内容，如果为空则从文件读取
+
+        Returns:
+            是否添加成功
+        """
+        if not content:
+            abs_path = self.kb_dir / file_path
+            if not abs_path.exists():
+                logger.warning(f"文件不存在: {abs_path}")
+                return False
+            content = self._read_file_content(abs_path)
+
+        if not content.strip():
+            return False
+
+        # 如果已有旧块，先移除
+        self.remove_document(file_path)
+
+        self._documents[file_path] = content
+
+        # 分块
+        text_chunks = self._split_text(content)
+        for idx, chunk_text in enumerate(text_chunks):
+            chunk_id = f"{file_path}#chunk_{idx}"
+            kc = KnowledgeChunk(
+                chunk_id=chunk_id,
+                file_path=file_path,
+                file_name=Path(file_path).name,
+                content=chunk_text,
+                chunk_index=idx,
+            )
+            self._chunks[chunk_id] = kc
+
+        logger.info(f"已添加文档到索引: {file_path} ({len(text_chunks)} 块)")
+        return True
+
+    def remove_document(self, file_path: str) -> int:
+        """
+        从索引中移除文档。
+
+        Args:
+            file_path: 文件相对路径
+
+        Returns:
+            移除的块数量
+        """
+        removed = 0
+        to_remove = [cid for cid in self._chunks if cid.startswith(f"{file_path}#")]
+        for cid in to_remove:
+            del self._chunks[cid]
+            removed += 1
+
+        self._documents.pop(file_path, None)
+        if removed > 0:
+            logger.info(f"已从索引移除文档: {file_path} ({removed} 块)")
+        return removed
+
+    # ==========================================================================
+    # 搜索
+    # ==========================================================================
+
+    def search(self, query: str, top_k: int = 5) -> List[KnowledgeChunk]:
+        """
+        搜索知识库，返回最相关的文本块。
+
+        Args:
+            query: 查询文本
+            top_k: 返回前 k 个结果
+
+        Returns:
+            按相关性排序的 KnowledgeChunk 列表
+        """
+        if not query or not self._chunks:
+            return []
+
+        # 构建文档列表
+        documents = [
+            (cid, chunk.content)
+            for cid, chunk in self._chunks.items()
+        ]
+
+        # 计算 TF-IDF 得分
+        scores = self._compute_tfidf(query, documents)
+
+        # 取 top_k
+        results = []
+        for chunk_id, score in list(scores.items())[:top_k]:
+            if chunk_id in self._chunks:
+                chunk = self._chunks[chunk_id]
+                chunk.score = score
+                results.append(chunk)
+
+        return results
+
+    def get_all_files(self) -> List[Dict[str, Any]]:
+        """
+        列出知识库中所有已索引的文件。
+
+        Returns:
+            文件信息列表 [{path, name, chunks, size}]
+        """
+        file_info: Dict[str, Dict] = {}
+        for cid, chunk in self._chunks.items():
+            fp = chunk.file_path
+            if fp not in file_info:
+                content = self._documents.get(fp, "")
+                file_info[fp] = {
+                    "path": fp,
+                    "name": chunk.file_name,
+                    "chunks": 0,
+                    "size": len(content),
+                }
+            file_info[fp]["chunks"] += 1
+
+        return list(file_info.values())
+
+    @property
+    def total_chunks(self) -> int:
+        """当前索引中的总块数"""
+        return len(self._chunks)
+
+    @property
+    def total_documents(self) -> int:
+        """当前索引中的总文档数"""
+        return len(self._documents)