250 lines
8.3 KiB
Python
250 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
InsightFlow LLM Client - Phase 4
|
||
用于与 Kimi API 交互，支持 RAG 问答和 Agent 功能
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
from collections.abc import AsyncGenerator
|
||
from dataclasses import dataclass
|
||
|
||
import httpx
|
||
|
||
KIMI_API_KEY = os.getenv("KIMI_API_KEY", "")
|
||
KIMI_BASE_URL = os.getenv("KIMI_BASE_URL", "https://api.kimi.com/coding")
|
||
|
||
|
||
@dataclass
|
||
class ChatMessage:
|
||
role: str
|
||
content: str
|
||
|
||
|
||
@dataclass
|
||
class EntityExtractionResult:
|
||
name: str
|
||
type: str
|
||
definition: str
|
||
confidence: float
|
||
|
||
|
||
@dataclass
|
||
class RelationExtractionResult:
|
||
source: str
|
||
target: str
|
||
type: str
|
||
confidence: float
|
||
|
||
|
||
class LLMClient:
|
||
"""Kimi API 客户端"""
|
||
|
||
def __init__(self, api_key: str = None, base_url: str = None):
|
||
self.api_key = api_key or KIMI_API_KEY
|
||
self.base_url = base_url or KIMI_BASE_URL
|
||
self.headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
|
||
|
||
async def chat(self, messages: list[ChatMessage], temperature: float = 0.3, stream: bool = False) -> str:
|
||
"""发送聊天请求"""
|
||
if not self.api_key:
|
||
raise ValueError("KIMI_API_KEY not set")
|
||
|
||
payload = {
|
||
"model": "k2p5",
|
||
"messages": [{"role": m.role, "content": m.content} for m in messages],
|
||
"temperature": temperature,
|
||
"stream": stream,
|
||
}
|
||
|
||
async with httpx.AsyncClient() as client:
|
||
response = await client.post(
|
||
f"{self.base_url}/v1/chat/completions", headers=self.headers, json=payload, timeout=120.0
|
||
)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
return result["choices"][0]["message"]["content"]
|
||
|
||
async def chat_stream(self, messages: list[ChatMessage], temperature: float = 0.3) -> AsyncGenerator[str, None]:
|
||
"""流式聊天请求"""
|
||
if not self.api_key:
|
||
raise ValueError("KIMI_API_KEY not set")
|
||
|
||
payload = {
|
||
"model": "k2p5",
|
||
"messages": [{"role": m.role, "content": m.content} for m in messages],
|
||
"temperature": temperature,
|
||
"stream": True,
|
||
}
|
||
|
||
async with httpx.AsyncClient() as client:
|
||
async with client.stream(
|
||
"POST", f"{self.base_url}/v1/chat/completions", headers=self.headers, json=payload, timeout=120.0
|
||
) as response:
|
||
response.raise_for_status()
|
||
async for line in response.aiter_lines():
|
||
if line.startswith("data: "):
|
||
data = line[6:]
|
||
if data == "[DONE]":
|
||
break
|
||
try:
|
||
chunk = json.loads(data)
|
||
delta = chunk["choices"][0]["delta"]
|
||
if "content" in delta:
|
||
yield delta["content"]
|
||
except (json.JSONDecodeError, KeyError, IndexError):
|
||
pass
|
||
|
||
async def extract_entities_with_confidence(
|
||
self, text: str
|
||
) -> tuple[list[EntityExtractionResult], list[RelationExtractionResult]]:
|
||
"""提取实体和关系，带置信度分数"""
|
||
prompt = f"""从以下会议文本中提取关键实体和它们之间的关系，以 JSON 格式返回：
|
||
|
||
文本：{text[:3000]}
|
||
|
||
要求：
|
||
1. entities: 每个实体包含 name(名称), type(类型: PROJECT/TECH/PERSON/ORG/OTHER), definition(一句话定义), confidence(置信度0-1)
|
||
2. relations: 每个关系包含 source(源实体名), target(目标实体名), type(关系类型: belongs_to/works_with/depends_on/mentions/related), confidence(置信度0-1)
|
||
3. 只返回 JSON 对象，格式: {{"entities": [...], "relations": [...]}}
|
||
|
||
示例：
|
||
{{
|
||
"entities": [
|
||
{{"name": "Project Alpha", "type": "PROJECT", "definition": "核心项目", "confidence": 0.95}},
|
||
{{"name": "K8s", "type": "TECH", "definition": "Kubernetes容器编排平台", "confidence": 0.88}}
|
||
],
|
||
"relations": [
|
||
{{"source": "Project Alpha", "target": "K8s", "type": "depends_on", "confidence": 0.82}}
|
||
]
|
||
}}"""
|
||
|
||
messages = [ChatMessage(role="user", content=prompt)]
|
||
content = await self.chat(messages, temperature=0.1)
|
||
|
||
import re
|
||
|
||
json_match = re.search(r"\{{.*?\}}", content, re.DOTALL)
|
||
if not json_match:
|
||
return [], []
|
||
|
||
try:
|
||
data = json.loads(json_match.group())
|
||
entities = [
|
||
EntityExtractionResult(
|
||
name=e["name"],
|
||
type=e.get("type", "OTHER"),
|
||
definition=e.get("definition", ""),
|
||
confidence=e.get("confidence", 0.8),
|
||
)
|
||
for e in data.get("entities", [])
|
||
]
|
||
relations = [
|
||
RelationExtractionResult(
|
||
source=r["source"],
|
||
target=r["target"],
|
||
type=r.get("type", "related"),
|
||
confidence=r.get("confidence", 0.8),
|
||
)
|
||
for r in data.get("relations", [])
|
||
]
|
||
return entities, relations
|
||
except Exception as e:
|
||
print(f"Parse extraction result failed: {e}")
|
||
return [], []
|
||
|
||
async def rag_query(self, query: str, context: str, project_context: dict) -> str:
|
||
"""RAG 问答 - 基于项目上下文回答问题"""
|
||
prompt = f"""你是一个专业的项目分析助手。基于以下项目信息回答问题：
|
||
|
||
## 项目信息
|
||
{json.dumps(project_context, ensure_ascii=False, indent=2)}
|
||
|
||
## 相关上下文
|
||
{context[:4000]}
|
||
|
||
## 用户问题
|
||
{query}
|
||
|
||
请用中文回答，保持简洁专业。如果信息不足，请明确说明。"""
|
||
|
||
messages = [
|
||
ChatMessage(role="system", content="你是一个专业的项目分析助手，擅长从会议记录中提取洞察。"),
|
||
ChatMessage(role="user", content=prompt),
|
||
]
|
||
|
||
return await self.chat(messages, temperature=0.3)
|
||
|
||
async def agent_command(self, command: str, project_context: dict) -> dict:
|
||
"""Agent 指令解析 - 将自然语言指令转换为结构化操作"""
|
||
prompt = f"""解析以下用户指令，转换为结构化操作：
|
||
|
||
## 项目信息
|
||
{json.dumps(project_context, ensure_ascii=False, indent=2)}
|
||
|
||
## 用户指令
|
||
{command}
|
||
|
||
请分析指令意图，返回 JSON 格式：
|
||
{{
|
||
"intent": "merge_entities|answer_question|edit_entity|create_relation|unknown",
|
||
"params": {{
|
||
// 根据 intent 不同，参数不同
|
||
}},
|
||
"explanation": "对用户指令的解释"
|
||
}}
|
||
|
||
意图说明：
|
||
- merge_entities: 合并实体，params 包含 source_names(源实体名列表), target_name(目标实体名)
|
||
- answer_question: 回答问题，params 包含 question(问题内容)
|
||
- edit_entity: 编辑实体，params 包含 entity_name(实体名), field(字段), value(新值)
|
||
- create_relation: 创建关系，params 包含 source(源实体), target(目标实体), relation_type(关系类型)
|
||
"""
|
||
|
||
messages = [ChatMessage(role="user", content=prompt)]
|
||
content = await self.chat(messages, temperature=0.1)
|
||
|
||
import re
|
||
|
||
json_match = re.search(r"\{{.*?\}}", content, re.DOTALL)
|
||
if not json_match:
|
||
return {"intent": "unknown", "explanation": "无法解析指令"}
|
||
|
||
try:
|
||
return json.loads(json_match.group())
|
||
except BaseException:
|
||
return {"intent": "unknown", "explanation": "解析失败"}
|
||
|
||
async def analyze_entity_evolution(self, entity_name: str, mentions: list[dict]) -> str:
|
||
"""分析实体在项目中的演变/态度变化"""
|
||
mentions_text = "\n".join(
|
||
[f"[{m.get('created_at', '未知时间')}] {m.get('text_snippet', '')}" for m in mentions[:20]] # 限制数量
|
||
)
|
||
|
||
prompt = f"""分析实体 "{entity_name}" 在项目中的演变和态度变化：
|
||
|
||
## 提及记录
|
||
{mentions_text}
|
||
|
||
请分析：
|
||
1. 该实体的角色/重要性变化
|
||
2. 相关方对它的态度变化
|
||
3. 关键时间节点
|
||
4. 总结性洞察
|
||
|
||
用中文回答，结构清晰。"""
|
||
|
||
messages = [ChatMessage(role="user", content=prompt)]
|
||
return await self.chat(messages, temperature=0.3)
|
||
|
||
|
||
# Singleton instance
|
||
_llm_client = None
|
||
|
||
|
||
def get_llm_client() -> LLMClient:
|
||
global _llm_client
|
||
if _llm_client is None:
|
||
_llm_client = LLMClient()
|
||
return _llm_client
|