细节优化
This commit is contained in:
Wenjie Zhang
parent
b2dbc17fec
commit
19a8e4714e
@ -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):
|
||||
|
||||
@ -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
|
||||
|
||||
return sorted_results
|
||||
@ -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')
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
return formatted_text
|
||||
@ -1,95 +1,95 @@
|
||||
<template>
|
||||
<div class="graph-container" ref="container"></div>
|
||||
</template>
|
||||
<div class="graph-container" ref="container"></div>
|
||||
</template>
|
||||
|
||||
<script setup>
|
||||
import { Graph } from "@antv/g6";
|
||||
import { onMounted, watch, ref } from 'vue';
|
||||
<script setup>
|
||||
import { Graph } from "@antv/g6";
|
||||
import { onMounted, watch, ref } from 'vue';
|
||||
|
||||
const props = defineProps({
|
||||
graphData: {
|
||||
type: Object,
|
||||
required: true,
|
||||
default: () => ({ nodes: [], edges: [] })
|
||||
}
|
||||
});
|
||||
|
||||
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 });
|
||||
</script>
|
||||
|
||||
<style scoped>
|
||||
.graph-container {
|
||||
background: #F7F7F7;
|
||||
border-radius: 16px;
|
||||
width: 100%;
|
||||
height: 600px;
|
||||
overflow: hidden;
|
||||
const props = defineProps({
|
||||
graphData: {
|
||||
type: Object,
|
||||
required: true,
|
||||
default: () => ({ nodes: [], edges: [] })
|
||||
}
|
||||
</style>
|
||||
});
|
||||
|
||||
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 });
|
||||
</script>
|
||||
|
||||
<style scoped>
|
||||
.graph-container {
|
||||
background: #F7F7F7;
|
||||
border-radius: 16px;
|
||||
width: 100%;
|
||||
height: 600px;
|
||||
overflow: hidden;
|
||||
}
|
||||
</style>
|
||||
@ -11,7 +11,7 @@
|
||||
<div class="graph-container layout-container" v-else>
|
||||
<HeaderComponent
|
||||
title="图数据库"
|
||||
:description="`${graphInfo?.database_name || ''} - 共 ${graphInfo?.entity_count || 0} 实体，${graphInfo?.relationship_count || 0} 个关系`"
|
||||
:description="`${graphInfo?.database_name || ''} - 共 ${graphInfo?.entity_count || 0} 实体，${graphInfo?.relationship_count || 0} 个关系。向量模型：${graphInfo?.embed_model_name || '未上传文件'}`"
|
||||
>
|
||||
<template #actions>
|
||||
<div class="status-wrapper">
|
||||
@ -43,7 +43,7 @@
|
||||
</div>
|
||||
</div>
|
||||
<div class="main" id="container" ref="container" v-show="graphData.nodes.length > 0"></div>
|
||||
<a-empty v-show="graphData.nodes.length === 0" style="padding: 4rem 0;"/>
|
||||
<a-empty v-show="graphData.nodes.length === 0" style="padding: 4rem 0;"/>
|
||||
|
||||
<a-modal
|
||||
:open="state.showModal" title="上传文件"
|
||||
@ -51,6 +51,11 @@
|
||||
@cancel="() => state.showModal = false"
|
||||
ok-text="添加到图数据库" cancel-text="取消"
|
||||
:confirm-loading="state.precessing">
|
||||
<div v-if="graphInfo?.embed_model_name">
|
||||
<p>当前图数据库向量模型：{{ graphInfo?.embed_model_name }}</p>
|
||||
<p>当前所选择的向量模型是 {{ configStore.config.embed_model }}</p>
|
||||
</div>
|
||||
<p v-else>第一次创建之后将无法修改向量模型，当前向量模型 {{ configStore.config.embed_model }}</p>
|
||||
<div class="upload">
|
||||
<a-upload-dragger
|
||||
class="upload-dragger"
|
||||
@ -58,7 +63,7 @@
|
||||
name="file"
|
||||
:fileList="fileList"
|
||||
:max-count="1"
|
||||
:disabled="state.precessing"
|
||||
:disabled="state.precessing || (graphInfo?.embed_model_name && graphInfo?.embed_model_name !== configStore.config.embed_model)"
|
||||
action="/api/data/upload"
|
||||
@change="handleFileUpload"
|
||||
@drop="handleDrop"
|
||||
@ -181,7 +186,7 @@ const loadSampleNodes = () => {
|
||||
graphData.nodes = data.result.nodes
|
||||
graphData.edges = data.result.edges
|
||||
console.log(graphData)
|
||||
randerGraph()
|
||||
setTimeout(() => randerGraph(), 500)
|
||||
})
|
||||
.catch((error) => {
|
||||
message.error(error.message);
|
||||
@ -195,12 +200,6 @@ const onSearch = () => {
|
||||
return
|
||||
}
|
||||
|
||||
const cur_embed_model = configStore.config.embed_model
|
||||
if (cur_embed_model !== 'zhipu-embedding-3') {
|
||||
message.error('当前不支持实体检索，请在设置中选择向量模型为 zhipu-embedding-3')
|
||||
return
|
||||
}
|
||||
|
||||
state.searchLoading = true
|
||||
fetch(`/api/data/graph/node?entity_name=${state.searchInput}`)
|
||||
.then((res) => {
|
||||
|
||||
@ -24,7 +24,7 @@
|
||||
<div class="section">
|
||||
<div class="card">
|
||||
<span class="label">{{ items?.embed_model.des }}</span>
|
||||
<a-select style="width: 200px"
|
||||
<a-select style="width: 300px"
|
||||
:value="configStore.config?.embed_model"
|
||||
@change="handleChange('embed_model', $event)"
|
||||
>
|
||||
@ -36,7 +36,7 @@
|
||||
</div>
|
||||
<div class="card">
|
||||
<span class="label">{{ items?.reranker.des }}</span>
|
||||
<a-select style="width: 200px"
|
||||
<a-select style="width: 300px"
|
||||
:value="configStore.config?.reranker"
|
||||
@change="handleChange('reranker', $event)"
|
||||
:disabled="!configStore.config.enable_reranker"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user