PM-pinou
65 lines
3.4 KiB
Python
65 lines
3.4 KiB
Python
"""Test full LLM orchestrator with actual tools"""
|
|
import sys, os, json, asyncio
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
|
|
import django; django.setup()
|
|
from analysis.session_store import SessionStore
|
|
from analysis.data_loader import load_csv_directory
|
|
from analysis.tool_registry import handle_call
|
|
import urllib.request
|
|
|
|
store = SessionStore()
|
|
store.drop_all()
|
|
csv = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'data', 'complex_test.csv')
|
|
lf, schema, rc, fc, mem = load_csv_directory(csv)
|
|
store.store_dataset('ds', lf, schema=schema, metadata={'row_count': rc, 'csv_glob': csv})
|
|
|
|
TOOLS = [{"type": "function", "function": {
|
|
"name": "profile_data", "description": "Get column statistics",
|
|
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}
|
|
}}, {"type": "function", "function": {
|
|
"name": "build_entity_profiles", "description": "Auto-detect entity and aggregate",
|
|
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "auto_detect": {"type": "boolean"}}, "required": ["dataset_id"]}
|
|
}}, {"type": "function", "function": {
|
|
"name": "run_clustering", "description": "Run HDBSCAN clustering",
|
|
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "cluster_columns": {"type": "array", "items": {"type": "string"}}, "algorithm": {"type": "string", "enum": ["hdbscan", "kmeans"]}}, "required": ["dataset_id", "cluster_columns"]}
|
|
}}, {"type": "function", "function": {
|
|
"name": "extract_features", "description": "Extract distinguishing features",
|
|
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "cluster_result_id": {"type": "string"}, "method": {"type": "string", "enum": ["zscore", "anova"]}}, "required": ["dataset_id", "cluster_result_id"]}
|
|
}}]
|
|
|
|
def call_llm(messages):
|
|
payload = json.dumps({"model": "deepseek-v4-flash", "messages": messages, "tools": TOOLS,
|
|
"tool_choice": "auto", "max_tokens": 4096}).encode()
|
|
req = urllib.request.Request('https://api.deepseek.com/chat/completions', data=payload,
|
|
headers={'Content-Type': 'application/json',
|
|
'Authorization': 'Bearer sk-360ef76d59674d6b8bc2eb160327dd39'},
|
|
method='POST')
|
|
return json.loads(urllib.request.urlopen(req, timeout=30).read())
|
|
|
|
msgs = [{"role": "user", "content": "Analyze dataset ds. Profile it, build entity profiles, cluster, extract features."}]
|
|
|
|
for step in range(10):
|
|
print(f'\nStep {step + 1}:')
|
|
response = call_llm(msgs)
|
|
msg = response["choices"][0]["message"]
|
|
|
|
if "tool_calls" not in msg or not msg["tool_calls"]:
|
|
print(f'Done: {msg.get("content", "")[:200]}')
|
|
break
|
|
|
|
for tc in msg["tool_calls"]:
|
|
func = tc.get("function", {})
|
|
name = func.get("name", "")
|
|
try:
|
|
args = json.loads(func.get("arguments", "{}"))
|
|
except json.JSONDecodeError:
|
|
args = {}
|
|
print(f' -> {name}({args})')
|
|
|
|
result = asyncio.run(handle_call(name, args))
|
|
content = json.dumps(result, default=str, ensure_ascii=False)[:2000]
|
|
print(f' <- {name}: ok' if 'error' not in result else f' <- ERROR: {result["error"]}')
|
|
|
|
msgs.append({"role": "tool", "tool_call_id": tc["id"], "content": content})
|