Improve search

core/messaging/ai.py

@@ -1,20 +1,68 @@
+import time
+
+from asgiref.sync import sync_to_async
+from django.utils import timezone
 from openai import AsyncOpenAI
 
-from core.models import AI
+from core.models import AI, AIRunLog
+
+
+def _prompt_metrics(prompt: list[dict]) -> tuple[int, int]:
+    rows = list(prompt or [])
+    message_count = len(rows)
+    prompt_chars = 0
+    for row in rows:
+        content = ""
+        if isinstance(row, dict):
+            content = str(row.get("content") or "")
+        else:
+            content = str(row or "")
+        prompt_chars += len(content)
+    return message_count, prompt_chars
 
 
 async def run_prompt(
-    prompt: list[str],
+    prompt: list[dict],
     ai: AI,
+    operation: str = "",
 ):
+    started_monotonic = time.perf_counter()
+    message_count, prompt_chars = _prompt_metrics(prompt)
+    run_log = await sync_to_async(AIRunLog.objects.create)(
+        user=ai.user,
+        ai=ai,
+        operation=str(operation or "").strip(),
+        model=str(ai.model or ""),
+        base_url=str(ai.base_url or ""),
+        status="running",
+        message_count=message_count,
+        prompt_chars=prompt_chars,
+    )
     cast = {"api_key": ai.api_key}
     if ai.base_url is not None:
         cast["base_url"] = ai.base_url
     client = AsyncOpenAI(**cast)
-    response = await client.chat.completions.create(
-        model=ai.model,
-        messages=prompt,
-    )
-    content = response.choices[0].message.content
-
-    return content
+    try:
+        response = await client.chat.completions.create(
+            model=ai.model,
+            messages=prompt,
+        )
+        content = response.choices[0].message.content
+        duration_ms = int((time.perf_counter() - started_monotonic) * 1000)
+        await sync_to_async(AIRunLog.objects.filter(id=run_log.id).update)(
+            status="ok",
+            response_chars=len(str(content or "")),
+            finished_at=timezone.now(),
+            duration_ms=duration_ms,
+            error="",
+        )
+        return content
+    except Exception as exc:
+        duration_ms = int((time.perf_counter() - started_monotonic) * 1000)
+        await sync_to_async(AIRunLog.objects.filter(id=run_log.id).update)(
+            status="failed",
+            finished_at=timezone.now(),
+            duration_ms=duration_ms,
+            error=str(exc),
+        )
+        raise