39 lines
1.7 KiB
Diff
39 lines
1.7 KiB
Diff
--- server/knowledge_worker.py
|
|
|
|
+++ local/knowledge_worker.py
|
|
|
|
@@ -17,7 +17,8 @@
|
|
|
|
|
|
|
|
-def enhance(candidate, database, options=None):
|
|
+def enhance(candidate, database, options=None, *, tenant_id="default", task_id=""):
|
|
"""Optional extraction via configured answer model, using only redacted evidence."""
|
|
from model_gateway import call_outlet
|
|
+ from model_usage import UsageRecorder
|
|
from knowledge_models import resolve_model, safe_model_error
|
|
import httpx
|
|
@@ -38,7 +39,10 @@
|
|
|
|
|
|
async def run():
|
|
- async with httpx.AsyncClient() as client:
|
|
- return await call_outlet(client, outlet, [{"role": "user", "content": prompt}],
|
|
- deadline=deadline, max_tokens=4096, temperature=0.1)
|
|
+ async with UsageRecorder(database, tenant_id=tenant_id, task_id=task_id,
|
|
+ purpose="knowledge") as recorder:
|
|
+ async with httpx.AsyncClient() as client:
|
|
+ return await call_outlet(client, outlet, [{"role": "user", "content": prompt}],
|
|
+ deadline=deadline, max_tokens=4096, temperature=0.1,
|
|
+ usage_recorder=recorder)
|
|
try:
|
|
result = asyncio.run(run())
|
|
@@ -181,5 +185,6 @@
|
|
|
|
if options['engine'] == 'model':
|
|
try:
|
|
- candidate, input_chars, output_chars = enhance(candidate, self.store.database, options)
|
|
+ candidate, input_chars, output_chars = enhance(
|
|
+ candidate, self.store.database, options, tenant_id=job["tenant_id"], task_id=job["id"])
|
|
except ModelOutputError as exc:
|
|
candidate = original_draft(candidate, exc)
|