57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
"""Quick test for SemanticRouter similarity on colloquial queries."""
|
|
import asyncio
|
|
import os
|
|
import dotenv
|
|
|
|
dotenv.load_dotenv()
|
|
|
|
from agentkit.chat.semantic_router import SemanticRouter
|
|
from agentkit.memory.embedder import OpenAIEmbedder
|
|
from agentkit.skills.registry import SkillRegistry
|
|
from agentkit.skills.loader import SkillLoader
|
|
from agentkit.server.config import ServerConfig
|
|
|
|
config = ServerConfig.from_yaml("agentkit.yaml")
|
|
key = os.environ.get("DASHSCOPE_API_KEY", "")
|
|
# Set API key and base_url for the first provider that needs it
|
|
for name, pconf in config.llm_config.providers.items():
|
|
if not pconf.api_key and key:
|
|
pconf.api_key = key
|
|
if not pconf.base_url:
|
|
pconf.base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
|
break
|
|
|
|
provider = config.llm_config.providers.get("test") or list(config.llm_config.providers.values())[0]
|
|
print(f"Using provider: api_key_len={len(provider.api_key)}, base_url={provider.base_url}")
|
|
|
|
embedder = OpenAIEmbedder(
|
|
api_key=provider.api_key,
|
|
base_url=provider.base_url,
|
|
model="text-embedding-v3",
|
|
)
|
|
|
|
router = SemanticRouter(embedder=embedder, similarity_low=0.4)
|
|
sr = SkillRegistry()
|
|
loader = SkillLoader(sr)
|
|
skills = loader.load_from_directory("configs/skills")
|
|
print(f"Loaded {len(skills)} skills: {[s.name for s in skills]}")
|
|
|
|
asyncio.run(router.build_index(sr))
|
|
print(f"SemanticRouter index size: {router._index.size}")
|
|
|
|
queries = [
|
|
"帮我看看代码有没有问题",
|
|
"对手怎么样",
|
|
"帮我写点东西",
|
|
"这个引用对不对",
|
|
"最近市场行情怎么样",
|
|
"review一下这段代码",
|
|
"做个SEO优化",
|
|
"monitor一下系统状态",
|
|
"审查代码",
|
|
"分析竞品策略",
|
|
]
|
|
for q in queries:
|
|
result = asyncio.run(router.route(q))
|
|
print(f"{q:30s} -> skill={str(result.skill_name):25s} sim={result.similarity:.3f} conf={result.confidence}")
|