Skip to content

feat: add Gemini function-calling framework with safe tool registry #37

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
GEEKYFOCUS wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

feat: add Gemini function-calling framework with safe tool registry #37

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 5 additions & 2 deletions .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,16 @@ jobs:
pip install pytest flake8

- name: Run linting
run: flake8 main.py stellar.py safety tests/redteam study.py --max-line-length=120 --ignore=E501,W503
run: flake8 main.py stellar.py safety tests/redteam study.py tools/ --max-line-length=120 --ignore=E501,W503

@coderabbitai coderabbitai Bot Jul 21, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Lint the new tool-test files too.

tests/test_tools.py and tests/conftest.py are changed Python files but are absent from the flake8 target, so their style or Pyflakes failures can merge unnoticed. Add both paths (or lint tests/ comprehensively).

Suggested CI adjustment 
- run: flake8 main.py stellar.py safety tests/redteam study.py tools/ --max-line-length=120 --ignore=E501,W503
+ run: flake8 main.py stellar.py safety tests/redteam tests/test_tools.py tests/conftest.py study.py tools/ --max-line-length=120 --ignore=E501,W503

As per path instructions, “CI enforces flake8, so style violations fail the build.”

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
run: flake8 main.py stellar.py safety tests/redteam study.py tools/ --max-line-length=120 --ignore=E501,W503
run: flake8 main.py stellar.py safety tests/redteam tests/test_tools.py tests/conftest.py study.py tools/ --max-line-length=120 --ignore=E501,W503
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In @.github/workflows/ci.yml at line 31, Update the flake8 command in the CI
workflow to include the changed test files tests/test_tools.py and
tests/conftest.py, or lint the tests/ directory comprehensively. Preserve the
existing max-line-length and ignore settings so these files are checked under
the same CI rules.

Source: Path instructions


- name: Check syntax
run: python -m compileall -q main.py stellar.py safety tests/redteam study.py
run: python -m compileall -q main.py stellar.py safety tests/redteam study.py tools/

- name: Run offline safety and red-team tests
run: pytest -q tests/redteam

- name: Run study generator tests
run: pytest -q tests/test_study.py

- name: Run tool registry tests
run: pytest -q tests/test_tools.py
12 changes: 6 additions & 6 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,8 @@ The platform is composed of three services:

| Method | Route | Purpose |
|--------|-------|---------|
| `POST` | `/chat` | Start or continue a chat session |
| `POST` | `/chat` | Start or continue a chat session; supports Gemini function-calling for zakat, Stellar info, and course search |
| `DELETE` | `/chat/{chat_id}` | Delete a chat session |
| `POST` | `/study/generate` | Generate schema-validated quizzes and flashcards |
| `GET` | `/ping` | Health check |

## 🚀 Getting Started
Expand Down Expand Up @@ -72,10 +71,11 @@ The API runs at `http://localhost:8000` — interactive docs at `http://localhos

### Environment Variables

| Variable | Description |
|----------|-------------|
| `GEMINI_API_KEY` | Google Gemini API key |
| `SAFETY_PIPELINE_ENABLED` | Layered policy enforcement; defaults to `true` |
| Variable | Description | Default |
|----------|-------------|---------|
| `GEMINI_API_KEY` | Google Gemini API key | — |
| `SAFETY_PIPELINE_ENABLED` | Layered policy enforcement | `true` |
| `BACKEND_API_URL` | Base URL for the dnb-backend REST API (tool: `search_courses`) | `https://dnb-backend-api.onrender.com` |

### Content-safety testing

Expand Down
115 changes: 102 additions & 13 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,19 @@
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import google.generativeai as genai
from google.generativeai import protos
import json
import os
from dotenv import load_dotenv
import logging
from typing import List, Optional
from typing import Any, List, Optional
import uuid

from stellar import router as stellar_router
from safety import InputGate, OutputCheck, SafetyPipeline, load_policy
from study import router as study_router
from tools import MAX_TOOL_ROUNDS, ToolCallRecord, get_registry
from tools.registry import run_tool_handler

# Configure logging
logging.basicConfig(level=logging.INFO)
Expand Down Expand Up @@ -96,6 +99,7 @@ class ChatResponse(BaseModel):
chat_id: str
history: List[Message]
moderation: Optional[Moderation] = None
tool_calls: Optional[List[ToolCallRecord]] = None


def classify_for_safety(prompt: str, candidate_ids: List[str]):
Expand Down Expand Up @@ -127,6 +131,10 @@ def classify_for_safety(prompt: str, candidate_ids: List[str]):
InputGate(safety_policy, classify_for_safety), OutputCheck(safety_policy)
)

# Tool registry for Gemini function-calling
tool_registry = get_registry()
tool_declarations = tool_registry.declarations()


def get_safety_settings():
return [
Expand Down Expand Up @@ -155,8 +163,75 @@ async def ping():
return {"************** Ping pong ping pong *************"}


def _run_with_tools(chat_session, prompt: str) -> tuple[str, list[ToolCallRecord]]:
"""Run a prompt through the model with function-calling loop.

Returns (final_text, tool_calls_list). Sync — the tool handlers use
a thread-pool executor for per-tool timeouts.
"""
tool_calls: list[ToolCallRecord] = []
content: Any = prompt

for _round in range(MAX_TOOL_ROUNDS):
response = chat_session.send_message(
content,
generation_config={
"temperature": 0.7,
"top_p": 0.8,
"top_k": 40,
"max_output_tokens": 2048,
},
)
candidate = response.candidates[0]
part = candidate.content.parts[0]
Comment on lines +185 to +186

@coderabbitai coderabbitai Bot Jul 21, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Unguarded indexing into candidates[0]/parts[0].

If the model response has no candidates (e.g. safety-blocked) or a candidate with no parts, this raises IndexError, which is caught only by the outermost bare except Exception in /chat (line 325) and surfaced to the client as a raw error string. Add explicit guards so this fails with a clear, non-leaking message.

Also applies to: 227-227

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@main.py` around lines 185 - 186, Guard the response handling around
candidate.content.parts in the chat flow before indexing candidates[0] or
parts[0]. If either collection is empty, return a clear user-facing error
without exposing raw exception details; preserve the existing processing path
when both contain an item. Apply the same protection at the corresponding later
access around the alternate response handling.


if part.text:
return part.text, tool_calls

if part.function_call:
fc = part.function_call
args_dict = dict(fc.args) if fc.args else {}

tool = tool_registry.get(fc.name)
if tool is None:
result = {"error": f"Unknown tool: {fc.name}"}
else:
result = run_tool_handler(tool, args_dict)

tool_calls.append(ToolCallRecord(
tool_name=fc.name,
args=args_dict,
result=json.dumps(result),
))

content = protos.Content(
parts=[protos.Part(
function_response=protos.FunctionResponse(
name=fc.name,
response=result,
)
)],
role="user",
)

# Exhausted rounds — force a final text answer
response = chat_session.send_message(
"Please provide your best final answer based on the information available.",
generation_config={
"temperature": 0.7,
"top_p": 0.8,
"top_k": 40,
"max_output_tokens": 2048,
},
)
final_text = response.candidates[0].content.parts[0].text or ""
return final_text, tool_calls

Comment on lines +166 to +229

@coderabbitai coderabbitai Bot Jul 21, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Fully synchronous Gemini + tool-calling loop runs on the event-loop thread.

generate() (and the _run_with_tools loop it now drives) performs up to MAX_TOOL_ROUNDS sequential blocking chat_session.send_message calls, each potentially followed by a blocking tool wait (up to timeout_seconds, e.g. 15s for zakat), plus one final forced-answer call — all inside async def chat(...) without ever being awaited or offloaded (SafetyPipeline._complete calls generator(...) synchronously; the non-safety branch calls generate(...) directly too). In the worst case this can block the event loop for a large multiple of a single Gemini round trip, stalling every other concurrent request on that worker.

Wrap the blocking generate() call in await asyncio.to_thread(generate, safety_prompt) (or FastAPI's run_in_threadpool) so the event loop stays responsive while the tool-calling loop runs.

As per path instructions, "Flag blocking calls inside async endpoints (network calls should be awaited or offloaded)" for this FastAPI service.

Also applies to: 240-274

🧰 Tools
🪛 ast-grep (0.44.1)

[info] 203-203: use jsonify instead of json.dumps for JSON output
Context: json.dumps(result)
Note: [CWE-116] Improper Encoding or Escaping of Output.

(use-jsonify)

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@main.py` around lines 166 - 229, Offload the synchronous generate flow from
the async chat handler by replacing direct calls to generate, including the
SafetyPipeline._complete path and non-safety branch, with await
asyncio.to_thread(generate, safety_prompt) or the project’s threadpool
equivalent. Preserve prompt selection and response handling while ensuring the
blocking _run_with_tools loop never executes on the event-loop thread.

Source: Path instructions


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
tool_calls_recorded: list[ToolCallRecord] = []

try:
logger.info(f"Received chat request: {request.prompt[:100]}...")

Expand All @@ -167,25 +242,36 @@ def generate(safety_prompt: str) -> str:
logger.info(f"Creating new chat session: {chat_id}")
model = genai.GenerativeModel(
'gemini-2.5-flash-preview-05-20',
tools=tool_declarations if tool_declarations else None,
safety_settings=get_safety_settings()
)
active_chats[chat_id] = model.start_chat(history=[])

context = f"Additional context: {request.context}\n\n" if request.context else ""
full_prompt = f"{ISLAMIC_CONTEXT}\n{context}User question: {safety_prompt}"
logger.info("Sending message to chat...")
response = active_chats[chat_id].send_message(
full_prompt,
generation_config={
"temperature": 0.7,
"top_p": 0.8,
"top_k": 40,
"max_output_tokens": 2048,
}
)
if not response.text:

chat_session = active_chats[chat_id]
if tool_declarations:
final_text, calls = _run_with_tools(chat_session, full_prompt)
tool_calls_recorded.extend(calls)
else:
response = chat_session.send_message(
full_prompt,
generation_config={
"temperature": 0.7,
"top_p": 0.8,
"top_k": 40,
"max_output_tokens": 2048,
}
)
if not response.text:
raise HTTPException(status_code=500, detail="Empty response from AI model")
final_text = response.text

if not final_text:
raise HTTPException(status_code=500, detail="Empty response from AI model")
return response.text
return final_text

enabled = os.getenv("SAFETY_PIPELINE_ENABLED", "true").lower() not in {"0", "false", "off"}
if enabled:
Expand Down Expand Up @@ -222,15 +308,18 @@ def generate(safety_prompt: str) -> str:
logger.warning(f"Error processing message in history: {str(e)}")
continue

response_text = safety_result.text if safety_result else generated_text

logger.info("Chat response generated successfully")
return ChatResponse(
response=safety_result.text if safety_result else generated_text,
response=response_text,
chat_id=chat_id,
history=history,
moderation=Moderation(
category_id=safety_result.category_id,
action=safety_result.action,
) if safety_result and safety_result.category_id else None,
tool_calls=tool_calls_recorded or None,
)

except Exception as e:
Expand Down
Loading
Loading