PM-pinou
73 lines
3.1 KiB
Python
73 lines
3.1 KiB
Python
"""Test LLM basic chat and orchestrator"""
|
|
import sys, os, json, urllib.request
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# Test 1: Basic chat
|
|
print('=== Test 1: Basic chat ===')
|
|
payload = json.dumps({
|
|
'model': 'deepseek-chat',
|
|
'messages': [
|
|
{'role': 'system', 'content': 'You are a data analysis assistant.'},
|
|
{'role': 'user', 'content': 'Profile dataset ds_001.'}
|
|
],
|
|
'max_tokens': 200,
|
|
}).encode()
|
|
|
|
req = urllib.request.Request('https://api.deepseek.com/chat/completions', data=payload,
|
|
headers={
|
|
'Content-Type': 'application/json',
|
|
'Authorization': 'Bearer sk-360ef76d59674d6b8bc2eb160327dd39',
|
|
},
|
|
method='POST')
|
|
with urllib.request.urlopen(req, timeout=15) as resp:
|
|
data = json.loads(resp.read())
|
|
print('Basic chat OK')
|
|
print('Response:', data['choices'][0]['message']['content'][:200])
|
|
|
|
# Test 2: LLM orchestrator
|
|
print('\n=== Test 2: LLM orchestrator ===')
|
|
os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
|
|
import django; django.setup()
|
|
from tianxuan.llm_orchestrator import run_llm_pipeline, LLMConfig
|
|
from analysis.session_store import SessionStore
|
|
from analysis.data_loader import load_csv_directory
|
|
from analysis.entity_detector import detect_entity_column
|
|
import asyncio
|
|
from analysis.tool_registry import _handle_run_clustering, _handle_extract_features
|
|
|
|
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})
|
|
|
|
cfg = LLMConfig(
|
|
base_url='https://api.deepseek.com',
|
|
api_key='sk-360ef76d59674d6b8bc2eb160327dd39',
|
|
model='deepseek-chat',
|
|
)
|
|
|
|
result = run_llm_pipeline('ds', config=cfg)
|
|
print('Orchestrator result:', json.dumps(result, ensure_ascii=False, indent=2)[:500])
|
|
|
|
# Test 3: Manual step-by-step (fallback if LLM fails)
|
|
if 'error' in result:
|
|
print('\n=== Test 3: Manual pipeline (fallback) ===')
|
|
entity_col = detect_entity_column('ds')['recommended']
|
|
print(f'Entity col: {entity_col}')
|
|
from analysis.entity_aggregator import aggregate_by_entity
|
|
agg_lf, features = aggregate_by_entity(lf, entity_col, schema)
|
|
df_agg = agg_lf.collect()
|
|
print(f'Aggregated: {len(df_agg)} entities, {len(features)} features')
|
|
store.store_dataset('entity_data', agg_lf, schema=dict(zip(features, ['']*len(features))),
|
|
metadata={'row_count': len(df_agg), 'csv_glob': csv})
|
|
clust_result = asyncio.run(_handle_run_clustering(dataset_id='entity_data',
|
|
cluster_columns=[c for c in features if not c.startswith('_')][:8],
|
|
algorithm='hdbscan', params={}, random_state=42))
|
|
print(f'Clusters: {clust_result.get("n_clusters")}')
|
|
cid = clust_result.get('cluster_result_id', '')
|
|
if cid and 'error' not in clust_result:
|
|
feat_result = asyncio.run(_handle_extract_features(
|
|
dataset_id='entity_data', cluster_result_id=cid,
|
|
top_k=10, method='zscore', save_to_db=True))
|
|
print(f'Features: {feat_result.get("n_features")} db_saved={feat_result.get("db_saved")}')
|