feat: mini llm fixes

This commit is contained in:
Alex committed 2026-06-18 13:35:37 +01:00
1 parent 7b0f5d777b
commit f27444fe48
2 files changed
+27 -12

No files matched your search

+19 -4
View File
@@ -545,7 +545,6 @@ class GoogleLLM(BaseLLM):
**kwargs,
):
"""Generate content using Google AI API without streaming."""
client = genai.Client(api_key=self.api_key)
system_instruction = None
if formatting == "openai":
messages, system_instruction = self._clean_messages_google(messages)
@@ -562,7 +561,7 @@ class GoogleLLM(BaseLLM):
if response_schema:
config.response_schema = response_schema
config.response_mime_type = "application/json"
response = client.models.generate_content(
response = self.client.models.generate_content(
model=model,
contents=messages,
config=config,
@@ -585,7 +584,6 @@ class GoogleLLM(BaseLLM):
**kwargs,
):
"""Generate content using Google AI API with streaming."""
client = genai.Client(api_key=self.api_key)
system_instruction = None
if formatting == "openai":
messages, system_instruction = self._clean_messages_google(messages)
@@ -607,7 +605,24 @@ class GoogleLLM(BaseLLM):
config.response_mime_type = "application/json"
# Check if we have both tools and file attachments
response = client.models.generate_content_stream(
has_attachments = False
for message in messages:
parts = getattr(message, "parts", None) or []
for part in parts:
if hasattr(part, "file_data") and part.file_data is not None:
has_attachments = True
break
if has_attachments:
break
messages_summary = self._summarize_messages_for_log(messages)
logging.info(
"GoogleLLM: Starting stream generation. Model: %s, Messages: %s, Has attachments: %s",
model,
messages_summary,
has_attachments,
)
response = self.client.models.generate_content_stream(
model=model,
contents=messages,
config=config,
+8 -8
View File
@@ -1076,8 +1076,6 @@ class OpenAILLM(BaseLLM):
Returns:
str: OpenAI file_id for the uploaded file.
"""
import logging
if "openai_file_id" in attachment:
return attachment["openai_file_id"]
file_path = attachment.get("path")
@@ -1085,12 +1083,14 @@ class OpenAILLM(BaseLLM):
if not self.storage.file_exists(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
try:
file_id = self.storage.process_file(
file_path,
lambda local_path, **kwargs: self.client.files.create(
file=open(local_path, "rb"), purpose="assistants"
).id,
)
def _upload(local_path, **_kwargs):
with open(local_path, "rb") as uploaded_file:
return self.client.files.create(
file=uploaded_file,
purpose="assistants",
).id
file_id = self.storage.process_file(file_path, _upload)
# Cache the OpenAI file id on the attachment row so we don't
# re-upload the same blob on the next LLM call. Prefer the PG