Skip to content

Make it possible to run with Gemini #500

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+47 −21
Open

Make it possible to run with Gemini #500

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
336c244
Update llm.py
yaruchyo Mar 6, 2025
ac4b7da
Update generate_code.py
yaruchyo Mar 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 35 additions & 21 deletions backend/llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,38 +268,52 @@ async def stream_gemini_response(

# Extract image URLs from messages
image_urls = []
for content_part in messages[-1]["content"]: # type: ignore
if content_part["type"] == "image_url": # type: ignore
image_url = content_part["image_url"]["url"] # type: ignore
if image_url.startswith("data:"): # type: ignore
# Extract base64 data and mime type for data URLs
mime_type = image_url.split(";")[0].split(":")[1] # type: ignore
base64_data = image_url.split(",")[1] # type: ignore
image_urls = [{"mime_type": mime_type, "data": base64_data}] # type: ignore
else:
# Store regular URLs
image_urls = [{"uri": image_url}] # type: ignore
break # Exit after first image URL
for content_part in messages[1]["content"]: # type: ignore
if isinstance(content_part, dict):
if content_part["type"] == "image_url": # type: ignore
image_url = content_part["image_url"]["url"] # type: ignore
if image_url.startswith("data:"): # type: ignore
# Extract base64 data and mime type for data URLs
mime_type = image_url.split(";")[0].split(":")[1] # type: ignore
base64_data = image_url.split(",")[1] # type: ignore
image_urls = [{"mime_type": mime_type, "data": base64_data}] # type: ignore
else:
# Store regular URLs
image_urls = [{"uri": image_url}] # type: ignore
break # Exit after first image URL

chat_history = []
for msg in messages:
if isinstance(msg["content"], str):
chat_history.append(f"{msg['role']}: {msg['content']}")
elif isinstance(msg["content"], list):
for element in msg["content"]:
if element.get("type") == "text":
chat_history.append(f"{msg['role']}: {element['text']}") # Extract the actual text

client = genai.Client(api_key=api_key) # type: ignore

full_response = ""
# if image_urls:
async for response in client.aio.models.generate_content_stream( # type: ignore
model=model.value,
contents={
"parts": [
{"text": messages[0]["content"]}, # type: ignore
model=model.value,
contents={
"parts": [
{"text": "\n\n".join(chat_history)}, # type: ignore
types.Part.from_bytes( # type: ignore
data=base64.b64decode(image_urls[0]["data"]), # type: ignore
mime_type=image_urls[0]["mime_type"], # type: ignore
),
] # type: ignore
}, # type: ignore
config=types.GenerateContentConfig( # type: ignore
temperature=0, max_output_tokens=8192
),
]
},
config=types.GenerateContentConfig( # type: ignore
temperature=0, max_output_tokens=8192
),
): # type: ignore
if response.text: # type: ignore
full_response += response.text # type: ignore
await callback(response.text) # type: ignore

completion_time = time.time() - start_time

return {"duration": completion_time, "code": full_response}
12 changes: 12 additions & 0 deletions backend/routes/generate_code.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ class ExtractedParams:
input_mode: InputMode
should_generate_images: bool
openai_api_key: str | None
gemini_api_key: str | None
anthropic_api_key: str | None
openai_base_url: str | None
generation_type: Literal["create", "update"]
Expand Down Expand Up @@ -109,6 +110,9 @@ async def extract_params(
anthropic_api_key = get_from_settings_dialog_or_env(
params, "anthropicApiKey", ANTHROPIC_API_KEY
)
gemini_api_key = get_from_settings_dialog_or_env(
params, "geminiApiKey", GEMINI_API_KEY
)

# Base URL for OpenAI API
openai_base_url: str | None = None
Expand Down Expand Up @@ -136,6 +140,7 @@ async def extract_params(
should_generate_images=should_generate_images,
openai_api_key=openai_api_key,
anthropic_api_key=anthropic_api_key,
gemini_api_key=gemini_api_key,
openai_base_url=openai_base_url,
generation_type=generation_type,
)
Expand Down Expand Up @@ -196,6 +201,7 @@ async def send_message(
openai_api_key = extracted_params.openai_api_key
openai_base_url = extracted_params.openai_base_url
anthropic_api_key = extracted_params.anthropic_api_key
gemini_api_key = extracted_params.gemini_api_key
should_generate_images = extracted_params.should_generate_images
generation_type = extracted_params.generation_type

Expand Down Expand Up @@ -277,6 +283,12 @@ async def process_chunk(content: str, variantIndex: int):
claude_model,
Llm.CLAUDE_3_5_SONNET_2024_06_20,
]
elif gemini_api_key:
variant_models = [
Llm.GEMINI_2_0_FLASH_EXP,
Llm.GEMINI_2_0_FLASH_EXP,
]

else:
await throw_error(
"No OpenAI or Anthropic API key found. Please add the environment variable OPENAI_API_KEY or ANTHROPIC_API_KEY to backend/.env or in the settings dialog. If you add it to .env, make sure to restart the backend server."
Expand Down