astra-vision/tests/test_core.py
Jeuner 16786fd3fd feat: image generation via ComfyUI and PDF/image upload
Zwei neue, sauber getrennte Fähigkeiten, jede in ihrem eigenen Modul:

- astra/comfyui.py: async HTTP-Client für einen lokalen ComfyUI-Server
  (z-image-turbo-Workflow), kennt nichts von der Pipeline.
- astra/documents.py: PDF-Textextraktion via pypdf, keine Netzwerkzugriffe.
- astra/tools.py: verdrahtet generate_image als natives Ollama-Tool-Call —
  Qwen 3.5 unterstützt Tools und Vision bereits nativ laut `ollama show`.

Dafür wurde astra/services.py so erweitert, dass NativeOllamaService
Ollamas native tool_calls im Streaming-Response erkennt und als
ChatCompletionChunk-Deltas an Pipecats bereits vorhandene, generische
Function-Calling-Maschinerie (_process_context/run_function_calls)
weiterreicht — die musste dafür nicht angefasst werden. trim_messages in
core.py bewahrt jetzt Tool-Roundtrips und Bild-Anhänge vollständig statt
sie auf role/content zu reduzieren.

Neuer Upload-Button im UI (Bild oder PDF, während eines laufenden
Gesprächs): PDFs gehen als Text, Bilder als Base64 über Qwens Vision in
den Gesprächskontext ein. Generierte Bilder werden über /api/media/<id>
ausgeliefert und per Datenkanal im Transkript angezeigt.

Kompletter Function-Calling-Roundtrip end-to-end gegen echtes Ollama und
echtes ComfyUI verifiziert (Modell ruft generate_image korrekt auf,
Bild wird erzeugt und im media_store abgelegt).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LVgSHNHdRx3UNTBodFmhRA
2026-09-07 15:56:18 +02:00

81 lines
3.4 KiB
Python

import unittest
from astra.core import Settings, build_request, local_origin_allowed, trim_messages
class CoreTests(unittest.TestCase):
def test_thinking_is_disabled_in_every_native_request(self):
request = build_request(Settings(), [{"role": "user", "content": "Hallo"}])
self.assertIs(request["think"], False)
self.assertIs(request["stream"], True)
self.assertEqual(request["model"], "qwen3.5:latest")
self.assertNotIn("think", request["options"])
self.assertEqual(request["options"]["num_ctx"], 4096)
def test_history_stays_bounded_and_preserves_system_and_latest_user(self):
messages = [{"role": "system", "content": "Deutsch"}]
for i in range(40):
messages.extend(
[
{"role": "user", "content": f"Frage {i}"},
{"role": "assistant", "content": "Antwort " * 100},
]
)
messages.append({"role": "user", "content": "Neueste Frage"})
trimmed = trim_messages(messages, max_chars=5000)
self.assertEqual(trimmed[0], messages[0])
self.assertEqual(trimmed[1]["role"], "user")
self.assertEqual(trimmed[-1], messages[-1])
self.assertLessEqual(sum(len(m["content"]) for m in trimmed), 5000)
self.assertEqual(len(messages), 82)
def test_rejects_foreign_web_origins(self):
self.assertTrue(local_origin_allowed("http://localhost:7860"))
self.assertTrue(local_origin_allowed("http://127.0.0.1:7860"))
self.assertFalse(local_origin_allowed("https://evil.example"))
self.assertFalse(local_origin_allowed("http://localhost.evil.example:7860"))
self.assertFalse(local_origin_allowed("null"))
def test_huge_last_message_is_bounded(self):
trimmed = trim_messages(
[
{"role": "system", "content": "Deutsch"},
{"role": "user", "content": "x" * 20000},
],
max_chars=5000,
)
self.assertLessEqual(sum(len(m["content"]) for m in trimmed), 5000)
def test_tool_round_trip_survives_trimming(self):
messages = [
{"role": "system", "content": "Deutsch"},
{"role": "user", "content": "Zeig mir ein Bild von einer Katze."},
{
"role": "assistant",
"content": "",
"tool_calls": [{"function": {"name": "generate_image", "arguments": {}}}],
},
{"role": "tool", "content": '{"status": "ok"}', "tool_call_id": "call_1"},
{"role": "assistant", "content": "Fertig, schau mal!"},
]
trimmed = trim_messages(messages, max_chars=5000)
roles = [m["role"] for m in trimmed]
self.assertEqual(roles, ["system", "user", "assistant", "tool", "assistant"])
self.assertEqual(trimmed[2]["tool_calls"][0]["function"]["name"], "generate_image")
self.assertEqual(trimmed[3]["tool_call_id"], "call_1")
def test_image_attachment_survives_trimming(self):
messages = [
{"role": "system", "content": "Deutsch"},
{
"role": "user",
"content": '[Hochgeladenes Bild "katze.png"]',
"images": ["base64data"],
},
]
trimmed = trim_messages(messages, max_chars=5000)
self.assertEqual(trimmed[1]["images"], ["base64data"])
if __name__ == "__main__":
unittest.main()