From 19a8e4714ebbd2dbd1284cb1f1c49102bc8c5546 Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Sun, 23 Feb 2025 16:39:52 +0800 Subject: [PATCH] =?UTF-8?q?=E7=BB=86=E8=8A=82=E4=BC=98=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/graphbase.py | 5 +- src/models/ollama_embedding.py | 52 ++++---- src/routers/chat_router.py | 26 ++-- src/utils/web_search.py | 20 +-- web/src/components/GraphContainer.vue | 182 +++++++++++++------------- web/src/views/GraphView.vue | 19 ++- web/src/views/SettingView.vue | 4 +- 7 files changed, 158 insertions(+), 150 deletions(-) diff --git a/src/core/graphbase.py b/src/core/graphbase.py index b33e3fac..399ec448 100644 --- a/src/core/graphbase.py +++ b/src/core/graphbase.py @@ -27,6 +27,7 @@ class GraphDatabase: self.kgdb_name = kgdb_name assert embed_model, "embed_model=None" self.embed_model = embed_model + self.embed_model_name = None def start(self): uri = os.environ.get("NEO4J_URI", "bolt://localhost:7687") @@ -87,7 +88,8 @@ class GraphDatabase: "relationship_count": relationship_count, "triples_count": triples_count, "labels": labels, - "status": self.status + "status": self.status, + "embed_model_name": self.embed_model_name } with self.driver.session() as session: @@ -171,6 +173,7 @@ class GraphDatabase: def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'): self.status = "processing" kgdb_name = kgdb_name or 'neo4j' + self.embed_model_name = self.embed_model_name or self.config.embed_model self.use_database(kgdb_name) # 切换到指定数据库 def read_triples(file_path): diff --git a/src/models/ollama_embedding.py b/src/models/ollama_embedding.py index 1b099be9..0819a46b 100644 --- a/src/models/ollama_embedding.py +++ b/src/models/ollama_embedding.py @@ -2,18 +2,20 @@ import os import requests import numpy as np from typing import List, Union, Dict + +from src.models.embedding import RemoteEmbeddingModel from src.utils.logging_config import setup_logger logger = setup_logger("OllamaEmbedding") -class OllamaEmbedding: +class OllamaEmbedding(RemoteEmbeddingModel): """ 使用 Ollama API 进行文本嵌入的类 """ def __init__(self, model_info: Dict, config) -> None: """ 初始化 Ollama Embedding 模型 - + Args: model_info: 模型信息字典 config: 配置对象 @@ -28,10 +30,10 @@ class OllamaEmbedding: def _get_embedding(self, text: str) -> List[float]: """ 获取单个文本的嵌入向量 - + Args: text: 输入文本 - + Returns: 嵌入向量 """ @@ -50,10 +52,10 @@ class OllamaEmbedding: def predict(self, messages: List[str]) -> List[List[float]]: """ 批量获取文本嵌入向量 - + Args: messages: 文本列表 - + Returns: 嵌入向量列表 """ @@ -63,23 +65,23 @@ class OllamaEmbedding: for i in range(0, len(messages), batch_size): batch = messages[i:i + batch_size] logger.info(f"Processing batch {i//batch_size + 1}, size: {len(batch)}") - + batch_embeddings = [] for text in batch: embedding = self._get_embedding(text) batch_embeddings.append(embedding) - + embeddings.extend(batch_embeddings) - + return embeddings def encode(self, messages: Union[str, List[str]]) -> List[List[float]]: """ 编码文本 - + Args: messages: 单个文本或文本列表 - + Returns: 嵌入向量列表 """ @@ -90,10 +92,10 @@ class OllamaEmbedding: def encode_queries(self, queries: List[str]) -> List[List[float]]: """ 编码查询文本 - + Args: queries: 查询文本列表 - + Returns: 查询文本的嵌入向量列表 """ @@ -107,7 +109,7 @@ class OllamaReranker: def __init__(self, config) -> None: """ 初始化 Ollama Reranker - + Args: config: 配置对象 """ @@ -119,16 +121,16 @@ class OllamaReranker: def compute_score(self, query: str, passage: str) -> float: """ 计算查询和文本段落之间的相关性分数 - + Args: query: 查询文本 passage: 段落文本 - + Returns: 相关性分数 """ prompt = f"Query: {query}\nPassage: {passage}\nRate the relevance of the passage to the query on a scale of 0 to 1:" - + try: response = requests.post( f"{self.base_url}/api/generate", @@ -139,7 +141,7 @@ class OllamaReranker: } ) response.raise_for_status() - + # 提取生成的数字作为分数 result = response.json()["response"].strip() try: @@ -148,7 +150,7 @@ class OllamaReranker: except ValueError: logger.warning(f"Could not parse score from response: {result}") return 0.0 - + except Exception as e: logger.error(f"Error computing rerank score: {str(e)}") return 0.0 @@ -156,12 +158,12 @@ class OllamaReranker: def rerank(self, query: str, passages: List[str], top_n: int = None) -> List[Dict]: """ 重新排序文本段落 - + Args: query: 查询文本 passages: 段落文本列表 top_n: 返回前 n 个结果 - + Returns: 排序后的结果列表，每个元素包含索引和分数 """ @@ -169,11 +171,11 @@ class OllamaReranker: for i, passage in enumerate(passages): score = self.compute_score(query, passage) scores.append({"index": i, "score": score}) - + # 按分数降序排序 sorted_results = sorted(scores, key=lambda x: x["score"], reverse=True) - + if top_n: sorted_results = sorted_results[:top_n] - - return sorted_results \ No newline at end of file + + return sorted_results \ No newline at end of file diff --git a/src/routers/chat_router.py b/src/routers/chat_router.py index 711fd1f8..ceb86c9e 100644 --- a/src/routers/chat_router.py +++ b/src/routers/chat_router.py @@ -27,26 +27,29 @@ def chat_post( history_manager = HistoryManager(history) - def make_chunk(content=None, status=None, history=None, reasoning_content=None): + def make_chunk(content=None, **kwargs): return json.dumps({ "response": content, - "reasoning_response": reasoning_content, - "history": history, "model_name": startup.config.model_name, - "status": status, "meta": meta, + **kwargs }, ensure_ascii=False).encode('utf-8') + b"\n" def generate_response(): modified_query = query + refs = None # 处理知识库检索 if meta and meta.get("enable_retrieval"): chunk = make_chunk(status="searching") yield chunk - modified_query, refs = startup.retriever(modified_query, history_manager.messages, meta) - refs_pool[cur_res_id] = refs + try: + modified_query, refs = startup.retriever(modified_query, history_manager.messages, meta) + except Exception as e: + logger.error(f"Retriever error: {e}") + yield make_chunk(message=f"Retriever error: {e}", status="error") + return messages = history_manager.get_history_with_msg(modified_query, max_rounds=meta.get('history_round')) history_manager.add_user(query) # 注意这里使用原始查询 @@ -56,7 +59,7 @@ def chat_post( reasoning_content = "" for delta in startup.model.predict(messages, stream=True): if not delta.content and hasattr(delta, 'reasoning_content'): - reasoning_content += delta.reasoning_content + reasoning_content += delta.reasoning_content or "" chunk = make_chunk(reasoning_content=reasoning_content, status="reasoning") yield chunk continue @@ -67,14 +70,15 @@ def chat_post( else: content += delta.content or "" - chunk = make_chunk(content=content, - reasoning_content=reasoning_content, - status="loading", - history=history_manager.update_ai(content)) + chunk = make_chunk(content=content, status="loading") yield chunk logger.debug(f"Final response: {content}") logger.debug(f"Final reasoning response: {reasoning_content}") + yield make_chunk(content=content, + status="finished", + history=history_manager.update_ai(content), + refs=refs) return StreamingResponse(generate_response(), media_type='application/json') diff --git a/src/utils/web_search.py b/src/utils/web_search.py index 781e4553..e66ecf86 100644 --- a/src/utils/web_search.py +++ b/src/utils/web_search.py @@ -16,11 +16,11 @@ class WebSearcher: def search(self, query: str, max_results: int = 1) -> List[Dict]: """ 使用 Tavily 搜索相关内容 - + Args: query: 搜索查询 max_results: 最大返回结果数 - + Returns: 搜索结果列表 """ @@ -30,7 +30,7 @@ class WebSearcher: search_depth="basic", max_results=max_results ) - + # 提取需要的信息 formatted_results = [] for result in search_results['results'][:max_results]: @@ -40,9 +40,9 @@ class WebSearcher: 'url': result.get('url', ''), 'score': result.get('score', 0) }) - + return formatted_results - + except Exception as e: logger.error(f"Error during web search: {str(e)}") return [] @@ -50,20 +50,20 @@ class WebSearcher: def format_search_results(self, results: List[Dict]) -> str: """ 将搜索结果格式化为文本 - + Args: results: 搜索结果列表 - + Returns: 格式化后的文本 """ if not results: return "没有找到相关的网络搜索结果。" - + formatted_text = "以下是相关的网络搜索结果：\n\n" for i, result in enumerate(results, 1): formatted_text += f"{i}. {result['title']}\n" formatted_text += f" {result['content']}\n" formatted_text += f" 来源: {result['url']}\n\n" - - return formatted_text \ No newline at end of file + + return formatted_text \ No newline at end of file diff --git a/web/src/components/GraphContainer.vue b/web/src/components/GraphContainer.vue index 30324c20..41794cdb 100644 --- a/web/src/components/GraphContainer.vue +++ b/web/src/components/GraphContainer.vue @@ -1,95 +1,95 @@ +
+ - - - \ No newline at end of file +}); + +const container = ref(null); +let graphInstance = null; + +const initGraph = () => { + graphInstance = new Graph({ + container: container.value, + width: container.value.offsetWidth, + height: container.value.offsetHeight, + autoFit: true, + autoResize: true, + layout: { + type: 'd3-force', + preventOverlap: true, + kr: 20, + collide: { + strength: 1.0, + }, + }, + node: { + type: 'circle', + style: { + labelText: (d) => d.data.label, + size: 70, + }, + palette: { + field: 'label', + color: 'tableau', + }, + }, + edge: { + type: 'line', + style: { + labelText: (d) => d.data.label, + labelBackground: '#fff', + endArrow: true, + }, + }, + behaviors: ['drag-element', 'zoom-canvas', 'drag-canvas'], + }); +}; + +const renderGraph = () => { + if (!graphInstance) { + initGraph(); + } + + const formattedData = { + nodes: props.graphData.nodes.map(node => ({ + id: node.id, + data: { label: node.name } + })), + edges: props.graphData.edges.map(edge => ({ + source: edge.source_id, + target: edge.target_id, + data: { label: edge.type } + })) + }; + + graphInstance.setData(formattedData); + graphInstance.render(); +}; + +onMounted(() => { + renderGraph(); + window.addEventListener('resize', renderGraph); +}); + +watch(() => props.graphData, renderGraph, { deep: true }); + + + \ No newline at end of file diff --git a/web/src/views/GraphView.vue b/web/src/views/GraphView.vue index 861acf99..99aa0330 100644 --- a/web/src/views/GraphView.vue +++ b/web/src/views/GraphView.vue @@ -11,7 +11,7 @@