jarvis-ai-assistant 0.3.17__py3-none-any.whl → 0.3.19__py3-none-any.whl

jarvis/jarvis_rag/cli.py

@@ -20,6 +20,7 @@ from jarvis.jarvis_utils.config import (
     get_rag_use_bm25,
     get_rag_use_rerank,
 )
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 def is_likely_text_file(file_path: Path) -> bool:
@@ -70,7 +71,10 @@ class _CustomPlatformLLM(LLMInterface):
 
     def __init__(self, platform: BasePlatform):
         self.platform = platform
-        print(f"✅ 使用自定义LLM: 平台='{platform.platform_name()}', 模型='{platform.name()}'")
+        PrettyOutput.print(
+            f"使用自定义LLM: 平台='{platform.platform_name()}', 模型='{platform.name()}'",
+            OutputType.INFO,
+        )
 
     def generate(self, prompt: str, **kwargs) -> str:
         return self.platform.chat_until_success(prompt)
@@ -84,13 +88,13 @@ def _create_custom_llm(platform_name: str, model_name: str) -> Optional[LLMInter
         registry = PlatformRegistry.get_global_platform_registry()
         platform_instance = registry.create_platform(platform_name)
         if not platform_instance:
-            print(f"❌ 错误: 平台 '{platform_name}' 未找到。")
+            PrettyOutput.print(f"错误: 平台 '{platform_name}' 未找到。", OutputType.ERROR)
             return None
         platform_instance.set_model_name(model_name)
         platform_instance.set_suppress_output(True)
         return _CustomPlatformLLM(platform_instance)
     except Exception as e:
-        print(f"❌ 创建自定义LLM时出错: {e}")
+        PrettyOutput.print(f"创建自定义LLM时出错: {e}", OutputType.ERROR)
         return None
 
 
@@ -114,10 +118,10 @@ def _load_ragignore_spec() -> Tuple[Optional[pathspec.PathSpec], Optional[Path]]
             with open(ignore_file_to_use, "r", encoding="utf-8") as f:
                 patterns = f.read().splitlines()
             spec = pathspec.PathSpec.from_lines("gitwildmatch", patterns)
-            print(f"✅ 加载忽略规则: {ignore_file_to_use}")
+            PrettyOutput.print(f"加载忽略规则: {ignore_file_to_use}", OutputType.SUCCESS)
             return spec, project_root_path
         except Exception as e:
-            print(f"⚠️ 加载 {ignore_file_to_use.name} 文件失败: {e}")
+            PrettyOutput.print(f"加载 {ignore_file_to_use.name} 文件失败: {e}", OutputType.WARNING)
 
     return None, None
 
@@ -166,7 +170,7 @@ def add_documents(
                 continue
 
             if path.is_dir():
-                print(f"🔍 正在扫描目录: {path}")
+                PrettyOutput.print(f"正在扫描目录: {path}", OutputType.INFO)
                 for item in path.rglob("*"):
                     if item.is_file() and is_likely_text_file(item):
                         files_to_process.add(item)
@@ -174,10 +178,10 @@ def add_documents(
                 if is_likely_text_file(path):
                     files_to_process.add(path)
                 else:
-                    print(f"⚠️ 跳过可能的二进制文件: {path}")
+                    PrettyOutput.print(f"跳过可能的二进制文件: {path}", OutputType.WARNING)
 
     if not files_to_process:
-        print("⚠️ 在指定路径中未找到任何文本文件。")
+        PrettyOutput.print("在指定路径中未找到任何文本文件。", OutputType.WARNING)
         return
 
     # 使用 .ragignore 过滤文件
@@ -198,14 +202,14 @@ def add_documents(
 
         ignored_count = initial_count - len(retained_files)
         if ignored_count > 0:
-            print(f"ℹ️ 根据 .ragignore 规则过滤掉 {ignored_count} 个文件。")
+            PrettyOutput.print(f"根据 .ragignore 规则过滤掉 {ignored_count} 个文件。", OutputType.INFO)
         files_to_process = retained_files
 
     if not files_to_process:
-        print("⚠️ 所有找到的文本文件都被忽略规则过滤掉了。")
+        PrettyOutput.print("所有找到的文本文件都被忽略规则过滤掉了。", OutputType.WARNING)
         return
 
-    print(f"✅ 发现 {len(files_to_process)} 个独立文件待处理。")
+    PrettyOutput.print(f"发现 {len(files_to_process)} 个独立文件待处理。", OutputType.INFO)
 
     try:
         pipeline = JarvisRAGPipeline(
@@ -229,26 +233,26 @@ def add_documents(
                     loader = TextLoader(str(file_path), encoding="utf-8")
 
                 docs_batch.extend(loader.load())
-                print(f"✅ 已加载: {file_path} (文件 {i + 1}/{total_files})")
+                PrettyOutput.print(f"已加载: {file_path} (文件 {i + 1}/{total_files})", OutputType.INFO)
             except Exception as e:
-                print(f"⚠️ 加载失败 {file_path}: {e}")
+                PrettyOutput.print(f"加载失败 {file_path}: {e}", OutputType.WARNING)
 
             # 当批处理已满或是最后一个文件时处理批处理
             if docs_batch and (len(docs_batch) >= batch_size or (i + 1) == total_files):
-                print(f"⚙️ 正在处理批次，包含 {len(docs_batch)} 个文档...")
+                PrettyOutput.print(f"正在处理批次，包含 {len(docs_batch)} 个文档...", OutputType.INFO)
                 pipeline.add_documents(docs_batch)
                 total_docs_added += len(docs_batch)
-                print(f"✅ 成功添加 {len(docs_batch)} 个文档。")
+                PrettyOutput.print(f"成功添加 {len(docs_batch)} 个文档。", OutputType.SUCCESS)
                 docs_batch = []  # 清空批处理
 
         if total_docs_added == 0:
-            print("❌ 未能成功加载任何文档。")
+            PrettyOutput.print("未能成功加载任何文档。", OutputType.ERROR)
             raise typer.Exit(code=1)
 
-        print(f"✅ 成功将 {total_docs_added} 个文档的内容添加至集合 '{collection_name}'。")
+        PrettyOutput.print(f"成功将 {total_docs_added} 个文档的内容添加至集合 '{collection_name}'。", OutputType.SUCCESS)
 
     except Exception as e:
-        print(f"❌ 发生严重错误: {e}")
+        PrettyOutput.print(f"发生严重错误: {e}", OutputType.ERROR)
         raise typer.Exit(code=1)
 
 
@@ -273,7 +277,7 @@ def list_documents(
         results = collection.get()  # 获取集合中的所有项目
 
         if not results or not results["metadatas"]:
-            print("ℹ️ 知识库中没有找到任何文档。")
+            PrettyOutput.print("知识库中没有找到任何文档。", OutputType.INFO)
             return
 
         # 从元数据中提取唯一的源文件路径
@@ -285,15 +289,15 @@ def list_documents(
                     sources.add(source)
 
         if not sources:
-            print("ℹ️ 知识库中没有找到任何带有源信息的文档。")
+            PrettyOutput.print("知识库中没有找到任何带有源信息的文档。", OutputType.INFO)
             return
 
-        print(f"📚 知识库 '{collection_name}' 中共有 {len(sources)} 个独立文档:")
+        PrettyOutput.print(f"知识库 '{collection_name}' 中共有 {len(sources)} 个独立文档:", OutputType.INFO)
         for i, source in enumerate(sorted(list(sources)), 1):
-            print(f"  {i}. {source}")
+            PrettyOutput.print(f"  {i}. {source}", OutputType.INFO)
 
     except Exception as e:
-        print(f"❌ 发生错误: {e}")
+        PrettyOutput.print(f"发生错误: {e}", OutputType.ERROR)
         raise typer.Exit(code=1)
 
 
@@ -330,14 +334,14 @@ def retrieve(
             use_rerank=use_rerank,
         )
 
-        print(f"🤔 正在为问题检索文档: '{question}'")
+        PrettyOutput.print(f"正在为问题检索文档: '{question}'", OutputType.INFO)
         retrieved_docs = pipeline.retrieve_only(question, n_results=n_results)
 
         if not retrieved_docs:
-            print("ℹ️ 未找到相关文档。")
+            PrettyOutput.print("未找到相关文档。", OutputType.INFO)
             return
 
-        print(f"✅ 成功检索到 {len(retrieved_docs)} 个文档:")
+        PrettyOutput.print(f"成功检索到 {len(retrieved_docs)} 个文档:", OutputType.SUCCESS)
         from jarvis.jarvis_utils.globals import console
 
         for i, doc in enumerate(retrieved_docs, 1):
@@ -350,7 +354,7 @@ def retrieve(
             console.print(Markdown(f"```\n{content}\n```"))
 
     except Exception as e:
-        print(f"❌ 发生错误: {e}")
+        PrettyOutput.print(f"发生错误: {e}", OutputType.ERROR)
         raise typer.Exit(code=1)
 
 
@@ -385,7 +389,7 @@ def query(
 ):
     """查询RAG知识库并打印答案。"""
     if model and not platform:
-        print("❌ 错误: --model 需要指定 --platform。")
+        PrettyOutput.print("错误: --model 需要指定 --platform。", OutputType.ERROR)
         raise typer.Exit(code=1)
 
     try:
@@ -407,17 +411,17 @@ def query(
             use_rerank=use_rerank,
         )
 
-        print(f"🤔 正在查询: '{question}'")
+        PrettyOutput.print(f"正在查询: '{question}'", OutputType.INFO)
         answer = pipeline.query(question)
 
-        print("💬 答案:")
+        PrettyOutput.print("答案:", OutputType.INFO)
         # 我们仍然可以使用 rich.markdown.Markdown，因为 PrettyOutput 底层使用了 rich
         from jarvis.jarvis_utils.globals import console
 
         console.print(Markdown(answer))
 
     except Exception as e:
-        print(f"❌ 发生错误: {e}")
+        PrettyOutput.print(f"发生错误: {e}", OutputType.ERROR)
         raise typer.Exit(code=1)
 
 
@@ -432,7 +436,7 @@ except ImportError:
 
 def _check_rag_dependencies():
     if not _RAG_INSTALLED:
-        print("❌ RAG依赖项未安装。" "请运行 'pip install \"jarvis-ai-assistant[rag]\"' 来使用此命令。")
+        PrettyOutput.print("RAG依赖项未安装。请运行 'pip install \"jarvis-ai-assistant[rag]\"' 来使用此命令。", OutputType.ERROR)
         raise typer.Exit(code=1)
 
 

jarvis/jarvis_rag/embedding_manager.py

@@ -3,6 +3,7 @@ from typing import List, cast
 from langchain_huggingface import HuggingFaceEmbeddings
 
 from .cache import EmbeddingCache
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 class EmbeddingManager:
@@ -23,7 +24,7 @@ class EmbeddingManager:
         """
         self.model_name = model_name
 
-        print(f"🚀 初始化嵌入管理器, 模型: '{self.model_name}'...")
+        PrettyOutput.print(f"初始化嵌入管理器, 模型: '{self.model_name}'...", OutputType.INFO)
 
         # 缓存的salt是模型名称，以防止冲突
         self.cache = EmbeddingCache(cache_dir=cache_dir, salt=self.model_name)
@@ -42,8 +43,8 @@ class EmbeddingManager:
                 show_progress=True,
             )
         except Exception as e:
-            print(f"❌ 加载嵌入模型 '{self.model_name}' 时出错: {e}")
-            print("请确保您已安装 'sentence_transformers' 和 'torch'。")
+            PrettyOutput.print(f"加载嵌入模型 '{self.model_name}' 时出错: {e}", OutputType.ERROR)
+            PrettyOutput.print("请确保您已安装 'sentence_transformers' 和 'torch'。", OutputType.WARNING)
             raise
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
@@ -71,8 +72,9 @@ class EmbeddingManager:
 
         # 为不在缓存中的文本计算嵌入
         if texts_to_embed:
-            print(
-                f"🔎 缓存未命中。正在为 {len(texts_to_embed)}/{len(texts)} 个文档计算嵌入。"
+            PrettyOutput.print(
+                f"缓存未命中。正在为 {len(texts_to_embed)}/{len(texts)} 个文档计算嵌入。",
+                OutputType.INFO,
             )
             new_embeddings = self.model.embed_documents(texts_to_embed)
 
@@ -83,7 +85,10 @@ class EmbeddingManager:
             for i, embedding in zip(indices_to_embed, new_embeddings):
                 cached_embeddings[i] = embedding
         else:
-            print(f"✅ 缓存命中。所有 {len(texts)} 个文档的嵌入均从缓存中检索。")
+            PrettyOutput.print(
+                f"缓存命中。所有 {len(texts)} 个文档的嵌入均从缓存中检索。",
+                OutputType.SUCCESS,
+            )
 
         return cast(List[List[float]], cached_embeddings)
 

jarvis/jarvis_rag/llm_interface.py

@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
 from jarvis.jarvis_agent import Agent as JarvisAgent
 from jarvis.jarvis_platform.base import BasePlatform
 from jarvis.jarvis_platform.registry import PlatformRegistry
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 class LLMInterface(ABC):
@@ -41,7 +42,7 @@ class ToolAgent_LLM(LLMInterface):
         """
         初始化工具-代理 LLM 包装器。
         """
-        print("🤖 已初始化工具 Agent 作为最终应答者。")
+        PrettyOutput.print("已初始化工具 Agent 作为最终应答者。", OutputType.INFO)
         self.allowed_tools = ["read_code", "execute_script"]
         # 为代理提供一个通用的系统提示
         self.system_prompt = "You are a helpful assistant. Please answer the user's question based on the provided context. You can use tools to find more information if needed."
@@ -83,7 +84,7 @@ class ToolAgent_LLM(LLMInterface):
             return str(final_answer)
 
         except Exception as e:
-            print(f"❌ Agent 在执行过程中发生错误: {e}")
+            PrettyOutput.print(f"Agent 在执行过程中发生错误: {e}", OutputType.ERROR)
             return "错误: Agent 未能成功生成回答。"
 
 
@@ -102,9 +103,9 @@ class JarvisPlatform_LLM(LLMInterface):
             self.registry = PlatformRegistry.get_global_platform_registry()
             self.platform: BasePlatform = self.registry.get_normal_platform()
             self.platform.set_suppress_output(False)  # 确保模型没有控制台输出
-            print(f"🚀 已初始化 Jarvis 平台 LLM，模型: {self.platform.name()}")
+            PrettyOutput.print(f"已初始化 Jarvis 平台 LLM，模型: {self.platform.name()}", OutputType.INFO)
         except Exception as e:
-            print(f"❌ 初始化 Jarvis 平台 LLM 失败: {e}")
+            PrettyOutput.print(f"初始化 Jarvis 平台 LLM 失败: {e}", OutputType.ERROR)
             raise
 
     def generate(self, prompt: str, **kwargs) -> str:
@@ -122,5 +123,5 @@ class JarvisPlatform_LLM(LLMInterface):
             # 使用健壮的chat_until_success方法
             return self.platform.chat_until_success(prompt)
         except Exception as e:
-            print(f"❌ 调用 Jarvis 平台模型时发生错误: {e}")
+            PrettyOutput.print(f"调用 Jarvis 平台模型时发生错误: {e}", OutputType.ERROR)
             return "错误: 无法从本地LLM获取响应。"

jarvis/jarvis_rag/rag_pipeline.py

@@ -8,6 +8,7 @@ from .llm_interface import JarvisPlatform_LLM, LLMInterface, ToolAgent_LLM
 from .query_rewriter import QueryRewriter
 from .reranker import Reranker
 from .retriever import ChromaRetriever
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 from jarvis.jarvis_utils.config import (
     get_rag_embedding_model,
     get_rag_rerank_model,
@@ -74,7 +75,7 @@ class JarvisRAGPipeline:
         self._reranker: Optional[Reranker] = None
         self._query_rewriter: Optional[QueryRewriter] = None
 
-        print("✅ JarvisRAGPipeline 初始化成功 (模型按需加载).")
+        PrettyOutput.print("JarvisRAGPipeline 初始化成功 (模型按需加载).", OutputType.SUCCESS)
 
     def _get_embedding_manager(self) -> EmbeddingManager:
         if self._embedding_manager is None:
@@ -193,7 +194,7 @@ class JarvisRAGPipeline:
         # 2. 为每个重写的查询检索初始候选文档
         all_candidate_docs = []
         for q in rewritten_queries:
-            print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
+            PrettyOutput.print(f"正在为查询变体 '{q}' 进行混合检索...", OutputType.INFO)
             candidates = self._get_retriever().retrieve(
                 q, n_results=n_results * 2, use_bm25=self.use_bm25
             )
@@ -208,7 +209,7 @@ class JarvisRAGPipeline:
 
         # 3. 根据*原始*查询对统一的候选池进行重排
         if self.use_rerank:
-            print(f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）...")
+            PrettyOutput.print(f"正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）...", OutputType.INFO)
             retrieved_docs = self._get_reranker().rerank(
                 query_text, unique_candidate_docs, top_n=n_results
             )
@@ -229,15 +230,15 @@ class JarvisRAGPipeline:
             )
         )
         if sources:
-            print(f"📚 根据以下文档回答:")
+            PrettyOutput.print("根据以下文档回答:", OutputType.INFO)
             for source in sources:
-                print(f"  - {source}")
+                PrettyOutput.print(f"  - {source}", OutputType.INFO)
 
         # 4. 创建最终提示并生成答案
         # 我们使用原始的query_text作为给LLM的最终提示
         prompt = self._create_prompt(query_text, retrieved_docs)
 
-        print("🤖 正在从LLM生成答案...")
+        PrettyOutput.print("正在从LLM生成答案...", OutputType.INFO)
         answer = self.llm.generate(prompt)
 
         return answer
@@ -259,7 +260,7 @@ class JarvisRAGPipeline:
         # 2. 检索候选文档
         all_candidate_docs = []
         for q in rewritten_queries:
-            print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
+            PrettyOutput.print(f"正在为查询变体 '{q}' 进行混合检索...", OutputType.INFO)
             candidates = self._get_retriever().retrieve(
                 q, n_results=n_results * 2, use_bm25=self.use_bm25
             )
@@ -273,7 +274,7 @@ class JarvisRAGPipeline:
 
         # 3. 重排
         if self.use_rerank:
-            print(f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排...")
+            PrettyOutput.print(f"正在对 {len(unique_candidate_docs)} 个候选文档进行重排...", OutputType.INFO)
             retrieved_docs = self._get_reranker().rerank(
                 query_text, unique_candidate_docs, top_n=n_results
             )

jarvis/jarvis_rag/reranker.py

@@ -4,6 +4,7 @@ from langchain.docstore.document import Document
 from sentence_transformers.cross_encoder import (  # type: ignore
     CrossEncoder,
 )
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 class Reranker:
@@ -19,9 +20,9 @@ class Reranker:
         参数:
             model_name (str): 要使用的Cross-Encoder模型的名称。
         """
-        print(f"🔍 正在初始化重排模型: {model_name}...")
+        PrettyOutput.print(f"正在初始化重排模型: {model_name}...", OutputType.INFO)
         self.model = CrossEncoder(model_name)
-        print("✅ 重排模型初始化成功。")
+        PrettyOutput.print("重排模型初始化成功。", OutputType.SUCCESS)
 
     def rerank(
         self, query: str, documents: List[Document], top_n: int = 5

jarvis/jarvis_rag/retriever.py

@@ -8,6 +8,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from rank_bm25 import BM25Okapi  # type: ignore
 
 from .embedding_manager import EmbeddingManager
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 class ChromaRetriever:
@@ -39,7 +40,10 @@ class ChromaRetriever:
         self.collection = self.client.get_or_create_collection(
             name=self.collection_name
         )
-        print(f"✅ ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。")
+        PrettyOutput.print(
+            f"ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。",
+            OutputType.SUCCESS,
+        )
 
         # BM25索引设置
         self.bm25_index_path = os.path.join(self.db_path, f"{collection_name}_bm25.pkl")
@@ -48,24 +52,24 @@ class ChromaRetriever:
     def _load_or_initialize_bm25(self):
         """从磁盘加载BM25索引或初始化一个新索引。"""
         if os.path.exists(self.bm25_index_path):
-            print("🔍 正在加载现有的 BM25 索引...")
+            PrettyOutput.print("正在加载现有的 BM25 索引...", OutputType.INFO)
             with open(self.bm25_index_path, "rb") as f:
                 data = pickle.load(f)
                 self.bm25_corpus = data["corpus"]
                 self.bm25_index = BM25Okapi(self.bm25_corpus)
-            print("✅ BM25 索引加载成功。")
+            PrettyOutput.print("BM25 索引加载成功。", OutputType.SUCCESS)
         else:
-            print("⚠️ 未找到 BM25 索引，将初始化一个新的。")
+            PrettyOutput.print("未找到 BM25 索引，将初始化一个新的。", OutputType.WARNING)
             self.bm25_corpus = []
             self.bm25_index = None
 
     def _save_bm25_index(self):
         """将BM25索引保存到磁盘。"""
         if self.bm25_index:
-            print("💾 正在保存 BM25 索引...")
+            PrettyOutput.print("正在保存 BM25 索引...", OutputType.INFO)
             with open(self.bm25_index_path, "wb") as f:
                 pickle.dump({"corpus": self.bm25_corpus, "index": self.bm25_index}, f)
-            print("✅ BM25 索引保存成功。")
+            PrettyOutput.print("BM25 索引保存成功。", OutputType.SUCCESS)
 
     def add_documents(
         self, documents: List[Document], chunk_size=1000, chunk_overlap=100
@@ -78,7 +82,10 @@ class ChromaRetriever:
         )
         chunks = text_splitter.split_documents(documents)
 
-        print(f"📄 已将 {len(documents)} 个文档拆分为 {len(chunks)} 个块。")
+        PrettyOutput.print(
+            f"已将 {len(documents)} 个文档拆分为 {len(chunks)} 个块。",
+            OutputType.INFO,
+        )
 
         if not chunks:
             return
@@ -97,7 +104,10 @@ class ChromaRetriever:
             documents=chunk_texts,
             metadatas=cast(Any, metadatas),
         )
-        print(f"✅ 成功将 {len(chunks)} 个块添加到 ChromaDB 集合中。")
+        PrettyOutput.print(
+            f"成功将 {len(chunks)} 个块添加到 ChromaDB 集合中。",
+            OutputType.SUCCESS,
+        )
 
         # 更新并保存BM25索引
         tokenized_chunks = [doc.split() for doc in chunk_texts]