128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""测试辅助类（不被 pytest 收集为 test）"""
|
||
|
||
import hashlib
|
||
import json
|
||
from typing import Any, Dict, List, Optional, Union
|
||
|
||
|
||
class MockLLMProvider:
|
||
"""
|
||
用预先准备好的回复序列响应 complete() / complete_messages()。每次 pop 一个；用完后重复最后一个。
|
||
|
||
response 项支持两种形式（向后兼容）:
|
||
- str : 纯文本回复（老用法）
|
||
- Dict[str, Any] : 结构化回复，可含 tool_calls，用于测试 tool calling agent
|
||
形如 {"content": "...", "tool_calls": [{"id": ..., "function": {...}}]}
|
||
"""
|
||
|
||
def __init__(self, responses):
|
||
if isinstance(responses, (str, dict)):
|
||
responses = [responses]
|
||
self._responses: List[Union[str, Dict[str, Any]]] = list(responses)
|
||
self.calls: List[Dict[str, Any]] = []
|
||
self._idx = 0
|
||
# 兼容 coding 模块对 _llm_config.temperature 的 fallback 探测
|
||
self._llm_config = type("_FakeCfg", (), {"temperature": 0.1, "model": "mock"})()
|
||
|
||
def _next(self) -> Union[str, Dict[str, Any]]:
|
||
if self._idx < len(self._responses):
|
||
r = self._responses[self._idx]
|
||
self._idx += 1
|
||
else:
|
||
r = self._responses[-1]
|
||
return r
|
||
|
||
def _build_response(self, raw: Union[str, Dict[str, Any]]):
|
||
from hy_memory.agent.llm_provider import LLMResponse
|
||
if isinstance(raw, str):
|
||
return LLMResponse(content=raw, tokens_used=10)
|
||
# Dict
|
||
content = raw.get("content", "") or ""
|
||
tool_calls = raw.get("tool_calls")
|
||
# tool_calls 内的 function.arguments 既允许 dict 也允许 JSON string；统一成 OpenAI 风格 string
|
||
if tool_calls:
|
||
normalized = []
|
||
for tc in tool_calls:
|
||
fn = tc.get("function", {})
|
||
args = fn.get("arguments")
|
||
if isinstance(args, dict):
|
||
args = json.dumps(args, ensure_ascii=False)
|
||
normalized.append({
|
||
"id": tc.get("id", "") or "",
|
||
"type": tc.get("type", "function"),
|
||
"function": {
|
||
"name": fn.get("name", ""),
|
||
"arguments": args or "{}",
|
||
},
|
||
})
|
||
tool_calls = normalized
|
||
return LLMResponse(content=content, tokens_used=10, tool_calls=tool_calls)
|
||
|
||
async def complete(self, prompt, max_tokens=500, temperature=0.7, tools=None, tool_choice=None, **kwargs):
|
||
self.calls.append({
|
||
"method": "complete",
|
||
"prompt": prompt,
|
||
"max_tokens": max_tokens,
|
||
"temperature": temperature,
|
||
"tools": tools,
|
||
"tool_choice": tool_choice,
|
||
})
|
||
return self._build_response(self._next())
|
||
|
||
async def complete_messages(self, messages, max_tokens=500, temperature=0.7,
|
||
tools=None, tool_choice=None, **kwargs):
|
||
self.calls.append({
|
||
"method": "complete_messages",
|
||
"messages": messages,
|
||
"max_tokens": max_tokens,
|
||
"temperature": temperature,
|
||
"tools": tools,
|
||
"tool_choice": tool_choice,
|
||
})
|
||
return self._build_response(self._next())
|
||
|
||
|
||
class FailingLLMProvider:
|
||
"""每次 complete / complete_messages 都 raise，用于 fail-safe 路径测试"""
|
||
|
||
def __init__(self):
|
||
self.calls = 0
|
||
self._llm_config = type("_FakeCfg", (), {"temperature": 0.1, "model": "mock"})()
|
||
|
||
async def complete(self, **kwargs):
|
||
self.calls += 1
|
||
raise RuntimeError("mock LLM failure")
|
||
|
||
async def complete_messages(self, **kwargs):
|
||
self.calls += 1
|
||
raise RuntimeError("mock LLM failure")
|
||
|
||
|
||
class MockEmbedService:
|
||
async def embed(self, text):
|
||
return self._hash_to_vec(text)
|
||
|
||
async def embed_batch(self, texts):
|
||
return [self._hash_to_vec(t) for t in texts]
|
||
|
||
async def embed_with_meta(self, text, use_cache: bool = True):
|
||
return self._hash_to_vec(text), {"cache_hit": False, "mock": True}
|
||
|
||
async def embed_batch_with_meta(self, texts, use_cache: bool = True):
|
||
vecs = [self._hash_to_vec(t) for t in texts]
|
||
metas = [{"cache_hit": False, "mock": True} for _ in texts]
|
||
return vecs, metas
|
||
|
||
async def embed_queued(self, text, use_cache: bool = True):
|
||
return self._hash_to_vec(text)
|
||
|
||
async def embed_queued_with_meta(self, text, use_cache: bool = True):
|
||
return self._hash_to_vec(text), {"cache_hit": False, "mock": True}
|
||
|
||
@staticmethod
|
||
def _hash_to_vec(t: str, dims: int = 1536):
|
||
h = hashlib.md5(t.encode()).digest() # 16 字节
|
||
v = [b / 255.0 for b in h] * (dims // 16)
|
||
return v[:dims]
|