From e7558cc34981aefb980937a00f27a56ea89a10f3 Mon Sep 17 00:00:00 2001 From: Geekyfocus <137049234+GEEKYFOCUS@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:18:25 +0100 Subject: [PATCH] feat: add Gemini function-calling framework with safe tool registry Implements a function-calling (tool-calling) system for the /chat endpoint, allowing the model to invoke platform capabilities via a pydantic-validated tool registry. Key changes: - New tools/ package with Tool, ToolRegistry, schema bridging - Pydantic v2 model_json_schema() -> Gemini FunctionDeclaration - Execution loop in /chat (max 4 rounds, per-tool timeouts via ThreadPoolExecutor) - ChatResponse.tool_calls optional field for frontend transparency - Three initial tools: calculate_zakat, get_stellar_info, search_courses - Stellar zakat logic refactored into tools/handlers.py for sharing between REST endpoint and tool - BACKEND_API_URL env var for course-search backend - 14 offline tests with fake model client (scripted function_call parts) - Updated CI to lint tools/ and run tool registry tests Closes #25 --- .github/workflows/ci.yml | 7 +- README.md | 12 +- main.py | 115 ++++++++++-- stellar.py | 153 +++++---------- tests/conftest.py | 4 + tests/test_tools.py | 388 +++++++++++++++++++++++++++++++++++++++ tools/__init__.py | 2 + tools/handlers.py | 226 +++++++++++++++++++++++ tools/registry.py | 139 ++++++++++++++ 9 files changed, 914 insertions(+), 132 deletions(-) create mode 100644 tests/test_tools.py create mode 100644 tools/__init__.py create mode 100644 tools/handlers.py create mode 100644 tools/registry.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a5e5cc1..d10ab58 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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 - 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 diff --git a/README.md b/README.md index cb9d631..b681147 100644 --- a/README.md +++ b/README.md @@ -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 @@ -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 diff --git a/main.py b/main.py index 1496c81..c140ca1 100644 --- a/main.py +++ b/main.py @@ -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) @@ -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]): @@ -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 [ @@ -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] + + 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 + + @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]}...") @@ -167,6 +242,7 @@ 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=[]) @@ -174,18 +250,28 @@ def generate(safety_prompt: str) -> str: 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: @@ -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: diff --git a/stellar.py b/stellar.py index e8b1591..716cf63 100644 --- a/stellar.py +++ b/stellar.py @@ -7,59 +7,29 @@ Strictly read-only: only public keys ever reach this service. Secret keys are never accepted, stored, or logged. + +Pure computation lives in tools/handlers.py so the Gemini function-calling +tools and the REST endpoints share the same logic. """ import logging -import os from decimal import Decimal from typing import Optional from fastapi import APIRouter, HTTPException from pydantic import BaseModel, Field -from stellar_sdk import Server -from stellar_sdk.exceptions import NotFoundError from stellar_sdk.strkey import StrKey -logger = logging.getLogger(__name__) - -router = APIRouter(tags=["stellar"]) - -# Network configuration — must match the rest of the platform -# (dnb-backend uses the same issuers in src/services/stellar/stellarService.js) -STELLAR_NETWORK = os.getenv("STELLAR_NETWORK", "testnet") - -HORIZON_URLS = { - "testnet": "https://horizon-testnet.stellar.org", - "public": "https://horizon.stellar.org", -} - -USDC_ISSUERS = { - # Circle's official USDC issuer on mainnet - "public": "GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN", - # Test USDC issuer used across the Deen Bridge platform on testnet - "testnet": "GBBD47IF6LWK7P7MDEVSCWR7DPUWV3NY3DTQEVFL4NAT4AQH3ZLLFLA5", -} - -# Zakat is 2.5% of zakatable wealth held for a lunar year, due once the -# total meets the nisab threshold. The nisab derives from the value of -# 85g of gold (or 595g of silver) and changes with market prices, so it -# is configurable; consult a scholar for rulings. -ZAKAT_RATE = Decimal("0.025") -DEFAULT_NISAB_USD = Decimal(os.getenv("ZAKAT_NISAB_USD", "6000")) - -DISCLAIMER = ( - "This is an automated estimate based on your on-chain USDC balance only. " - "Zakat rulings depend on your full wealth, debts, and the hawl (lunar year). " - "Please consult a qualified scholar for a definitive ruling." +from tools.handlers import ( + DISCLAIMER, + _calculate_zakat, + _horizon_url, + _usdc_issuer, ) +logger = logging.getLogger(__name__) -def horizon_url() -> str: - return HORIZON_URLS.get(STELLAR_NETWORK, HORIZON_URLS["testnet"]) - - -def usdc_issuer() -> str: - return USDC_ISSUERS.get(STELLAR_NETWORK, USDC_ISSUERS["testnet"]) +router = APIRouter(tags=["stellar"]) class ZakatRequest(BaseModel): @@ -81,35 +51,46 @@ class ZakatResponse(BaseModel): disclaimer: str -def fetch_usdc_balance(public_key: str) -> Optional[Decimal]: - """Return the account's USDC balance, or None if it has no USDC trustline. - - Raises HTTPException(404) if the account does not exist on this network. +def compute_zakat(public_key: str, nisab_usd: Optional[float] = None) -> ZakatResponse: + """Pure computation shared by the REST endpoint and the tool handler. + Returns a ZakatResponse; raises HTTPException on invalid key. """ - server = Server(horizon_url()) - try: - account = server.accounts().account_id(public_key).call() - except NotFoundError: + public_key = public_key.strip() + if not StrKey.is_valid_ed25519_public_key(public_key): raise HTTPException( - status_code=404, - detail=f"Account not found on the Stellar {STELLAR_NETWORK} network.", + status_code=400, + detail="Invalid Stellar public key. Expected a 56-character key starting with G.", ) - for balance in account.get("balances", []): - if ( - balance.get("asset_code") == "USDC" - and balance.get("asset_issuer") == usdc_issuer() - ): - return Decimal(balance["balance"]) - return None + + logger.info("Zakat lookup for %s", public_key[:8]) + result = _calculate_zakat(public_key, nisab_usd) + + if "error" in result: + raise HTTPException(status_code=400, detail=result["error"]) + + nisab_val = Decimal(result["nisab_usd"]) + zakat_due = Decimal(result["zakat_due"]) + + return ZakatResponse( + network=result["network"], + public_key=result["public_key"], + has_usdc_trustline=result["has_usdc_trustline"], + usdc_balance=result["usdc_balance"], + nisab_usd=str(nisab_val), + zakat_rate=result["zakat_rate"], + zakat_due=str(zakat_due), + message=result["message"], + disclaimer=result.get("disclaimer", DISCLAIMER), + ) @router.get("/stellar/info") async def stellar_info(): """Public configuration of this service's Stellar integration.""" return { - "network": STELLAR_NETWORK, - "horizon": horizon_url(), - "usdc_issuer": usdc_issuer(), + "network": _horizon_url(), + "horizon": _horizon_url(), + "usdc_issuer": _usdc_issuer(), "features": ["zakat"], } @@ -117,54 +98,4 @@ async def stellar_info(): @router.post("/zakat", response_model=ZakatResponse) async def calculate_zakat(request: ZakatRequest): """Calculate zakat due on a wallet's on-chain USDC balance.""" - public_key = request.public_key.strip() - if not StrKey.is_valid_ed25519_public_key(public_key): - raise HTTPException( - status_code=400, - detail="Invalid Stellar public key. Expected a 56-character key starting with G.", - ) - - logger.info("Zakat lookup for %s on %s", public_key[:8], STELLAR_NETWORK) - balance = fetch_usdc_balance(public_key) - nisab = Decimal(str(request.nisab_usd)) if request.nisab_usd else DEFAULT_NISAB_USD - - if balance is None: - return ZakatResponse( - network=STELLAR_NETWORK, - public_key=public_key, - has_usdc_trustline=False, - usdc_balance="0", - nisab_usd=str(nisab), - zakat_rate=str(ZAKAT_RATE), - zakat_due="0", - message=( - "This account has no USDC trustline, so it holds no USDC. " - "Add a USDC trustline in your wallet to hold USDC on Stellar." - ), - disclaimer=DISCLAIMER, - ) - - meets_nisab = balance >= nisab - zakat_due = (balance * ZAKAT_RATE).quantize(Decimal("0.0000001")) if meets_nisab else Decimal("0") - if meets_nisab: - message = ( - f"Your USDC balance of {balance} meets the nisab threshold of {nisab} USD. " - f"If held for a full lunar year, the zakat due is {zakat_due} USDC (2.5%)." - ) - else: - message = ( - f"Your USDC balance of {balance} is below the nisab threshold of {nisab} USD, " - "so no zakat is due on this balance alone." - ) - - return ZakatResponse( - network=STELLAR_NETWORK, - public_key=public_key, - has_usdc_trustline=True, - usdc_balance=str(balance), - nisab_usd=str(nisab), - zakat_rate=str(ZAKAT_RATE), - zakat_due=str(zakat_due), - message=message, - disclaimer=DISCLAIMER, - ) + return compute_zakat(request.public_key, request.nisab_usd) diff --git a/tests/conftest.py b/tests/conftest.py index 9a107c2..139c606 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,6 @@ """Test configuration for importing repository modules without packaging the app.""" +import os import sys from pathlib import Path @@ -7,3 +8,6 @@ ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) + +# Fake key so modules that check at import time (main.py) don't fail +os.environ.setdefault("GEMINI_API_KEY", "test-fake-key") diff --git a/tests/test_tools.py b/tests/test_tools.py new file mode 100644 index 0000000..470b2eb --- /dev/null +++ b/tests/test_tools.py @@ -0,0 +1,388 @@ +"""Tests for the Gemini function-calling framework — no live API calls. + +Uses a fake model client that emits scripted function_call / text parts. +""" + +import json +import time +from unittest.mock import MagicMock, patch + +import pytest +from pydantic import BaseModel, Field + +from tools.registry import ( + MAX_TOOL_ROUNDS, + Tool, + ToolCallRecord, + ToolRegistry, + _clean_schema, + get_tool_declarations, + run_tool_handler, +) + + +# --------------------------------------------------------------------------- +# Fake model response helpers +# --------------------------------------------------------------------------- + + +class FakePart: + def __init__(self, text=None, function_call=None): + self._text = text + self._function_call = function_call + + @property + def text(self): + return self._text + + @property + def function_call(self): + return self._function_call + + +class FakeCandidate: + def __init__(self, part): + self.content = MagicMock() + self.content.parts = [part] + + +class FakeResponse: + def __init__(self, part): + self.candidates = [FakeCandidate(part)] + + +class FakeFunctionCall: + def __init__(self, name, args): + self.name = name + self.args = args + + +# --------------------------------------------------------------------------- +# Schema bridging tests +# --------------------------------------------------------------------------- + + +class SampleArgs(BaseModel): + name: str = Field(..., description="A name") + count: int = Field(5, description="A count") + tags: list[str] = Field(default_factory=list, description="Tags") + + +def test_clean_schema_strips_unsupported_keys(): + raw = { + "$schema": "http://json-schema.org/draft-07/schema#", + "$defs": {"Foo": {"type": "string"}}, + "title": "Sample", + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + } + cleaned = _clean_schema(raw) + assert "$schema" not in cleaned + assert "$defs" not in cleaned + assert "title" not in cleaned + assert cleaned["type"] == "object" + assert cleaned["required"] == ["name"] + + +def test_get_tool_declarations_from_pydantic(): + tool = Tool( + name="test_tool", + description="A test tool", + args_schema=SampleArgs, + handler=lambda **kw: {"ok": True}, + timeout_seconds=5, + ) + decls = get_tool_declarations([tool]) + assert len(decls) == 1 + decl = decls[0] + assert decl["name"] == "test_tool" + assert "description" in decl + assert decl["parameters"]["type"] == "object" + assert "name" in decl["parameters"]["properties"] + assert "count" in decl["parameters"]["properties"] + assert "tags" in decl["parameters"]["properties"] + assert "name" in decl["parameters"]["required"] + # Unsupported keys stripped + assert "$schema" not in decl["parameters"] + assert "$defs" not in decl["parameters"] + + +# --------------------------------------------------------------------------- +# Tool registry tests +# --------------------------------------------------------------------------- + + +def test_registry_register_and_get(): + registry = ToolRegistry() + tool = Tool( + name="hello", + description="Says hello", + args_schema=SampleArgs, + handler=lambda **kw: {"msg": "hello"}, + timeout_seconds=1, + ) + registry.register(tool) + assert registry.get("hello") is tool + assert registry.get("unknown") is None + + +def test_registry_duplicate_raises(): + registry = ToolRegistry() + tool = Tool( + name="dup", + description="", + args_schema=SampleArgs, + handler=lambda **kw: {}, + timeout_seconds=1, + ) + registry.register(tool) + with pytest.raises(ValueError, match="already registered"): + registry.register(tool) + + +def test_registry_all_and_declarations(): + registry = ToolRegistry() + t1 = Tool(name="a", description="", args_schema=SampleArgs, handler=lambda **kw: {}, timeout_seconds=1) + t2 = Tool(name="b", description="", args_schema=SampleArgs, handler=lambda **kw: {}, timeout_seconds=1) + registry.register(t1) + registry.register(t2) + assert len(registry.all) == 2 + assert len(registry.declarations()) == 2 + + +# --------------------------------------------------------------------------- +# Arg validation tests +# --------------------------------------------------------------------------- + + +class StrictArgs(BaseModel): + name: str = Field(..., min_length=1) + age: int = Field(..., gt=0) + + +def test_tool_handler_validates_args(): + handler_called = [] + + def my_handler(name: str, age: int): + handler_called.append((name, age)) + return {"ok": True} + + tool = Tool( + name="strict", + description="", + args_schema=StrictArgs, + handler=my_handler, + timeout_seconds=5, + ) + + # Valid args + result = run_tool_handler(tool, {"name": "Alice", "age": 30}) + assert result == {"ok": True} + assert len(handler_called) == 1 + + # Invalid args — fails pydantic validation + handler_called.clear() + result = run_tool_handler(tool, {"name": "", "age": -1}) + assert "error" in result + + +# --------------------------------------------------------------------------- +# Timeout enforcement tests +# --------------------------------------------------------------------------- + + +def slow_handler(**kw): + time.sleep(10) + return {"done": True} + + +def test_tool_timeout_is_enforced(): + tool = Tool( + name="slow", + description="", + args_schema=SampleArgs, + handler=slow_handler, + timeout_seconds=1, + ) + started = time.time() + result = run_tool_handler(tool, {"name": "test", "count": 1}) + elapsed = time.time() - started + assert "error" in result + assert "timed out" in result["error"] + assert elapsed < 5 # Should not wait the full 10s + + +# --------------------------------------------------------------------------- +# Loop bound tests — fake chat that emits function_calls +# --------------------------------------------------------------------------- + + +def _make_fake_chat(responses, final_text="Final answer"): + """Create a fake chat session that yields scripted function_call responses.""" + chat = MagicMock() + reply_index = [0] + + def side_effect(content, **kw): + idx = reply_index[0] + reply_index[0] += 1 + if idx < len(responses): + fc_data = responses[idx] + fc = FakeFunctionCall(name=fc_data["name"], args=fc_data.get("args", {})) + return FakeResponse(FakePart(function_call=fc)) + return FakeResponse(FakePart(text=final_text)) + + chat.send_message = side_effect + return chat + + +@patch("tools.registry.run_tool_handler", return_value={"ok": True}) +def test_loop_terminates_on_text(mock_handler): + """When the model returns text (no function_call), loop exits immediately.""" + chat = _make_fake_chat([], "Direct text answer") + from main import _run_with_tools + + # Need a real registry for _run_with_tools to work + registry = ToolRegistry() + tool = Tool( + name="dummy", + description="", + args_schema=SampleArgs, + handler=lambda **kw: {}, + timeout_seconds=5, + ) + registry.register(tool) + # Temp patch the module-level registry + import main as main_module + original_reg = main_module.tool_registry + main_module.tool_registry = registry + original_decls = main_module.tool_declarations + main_module.tool_declarations = registry.declarations() + + try: + text, calls = _run_with_tools(chat, "test prompt") + assert text == "Direct text answer" + assert calls == [] + finally: + main_module.tool_registry = original_reg + main_module.tool_declarations = original_decls + + +@patch("tools.registry.run_tool_handler", return_value={"ok": True}) +def test_loop_bounded(mock_handler): + """Model that keeps returning function_calls cannot loop forever.""" + # Script MAX_TOOL_ROUNDS function_call responses — after that the forced + # text message will get the final_text from the mock. + responses = [{"name": "dummy", "args": {"name": "x", "count": 1}} for _ in range(MAX_TOOL_ROUNDS)] + chat = _make_fake_chat(responses, "Final after loop") + from main import _run_with_tools + + registry = ToolRegistry() + tool = Tool( + name="dummy", + description="", + args_schema=SampleArgs, + handler=lambda **kw: {}, + timeout_seconds=5, + ) + registry.register(tool) + import main as main_module + original_reg = main_module.tool_registry + main_module.tool_registry = registry + original_decls = main_module.tool_declarations + main_module.tool_declarations = registry.declarations() + + try: + text, calls = _run_with_tools(chat, "test") + # Should have MAX_TOOL_ROUNDS calls, then forced text + assert len(calls) == MAX_TOOL_ROUNDS + assert text == "Final after loop" + finally: + main_module.tool_registry = original_reg + main_module.tool_declarations = original_decls + + +# --------------------------------------------------------------------------- +# Tool execution tests +# --------------------------------------------------------------------------- + + +def test_tool_handler_called_with_validated_args(): + recorded = {} + + def my_handler(name: str, count: int = 5, **kwargs): + recorded["name"] = name + recorded["count"] = count + return {"result": f"Hello {name}"} + + tool = Tool( + name="greet", + description="Greets someone", + args_schema=SampleArgs, + handler=my_handler, + timeout_seconds=5, + ) + result = run_tool_handler(tool, {"name": "Alice", "count": 3}) + assert result == {"result": "Hello Alice"} + assert recorded == {"name": "Alice", "count": 3} + + +def test_tool_handler_defaults(): + recorded = {} + + def my_handler(name: str, count: int = 5, **kwargs): + recorded["count"] = count + return {} + + tool = Tool( + name="defaults", + description="", + args_schema=SampleArgs, + handler=my_handler, + timeout_seconds=5, + ) + run_tool_handler(tool, {"name": "Bob"}) + assert recorded["count"] == 5 + + +def test_tool_handler_exception_becomes_error(): + def broken(**kw): + raise ValueError("Something broke") + + tool = Tool( + name="broken", + description="", + args_schema=SampleArgs, + handler=broken, + timeout_seconds=5, + ) + result = run_tool_handler(tool, {"name": "test"}) + assert "error" in result + assert "Something broke" in result["error"] + + +# --------------------------------------------------------------------------- +# ToolCallRecord model +# --------------------------------------------------------------------------- + + +def test_tool_call_record_creation(): + rec = ToolCallRecord( + tool_name="zakat", + args={"public_key": "GABCDEF"}, + result='{"zakat_due": "10.5"}', + ) + assert rec.tool_name == "zakat" + assert rec.args["public_key"] == "GABCDEF" + assert json.loads(rec.result)["zakat_due"] == "10.5" + + +# --------------------------------------------------------------------------- +# Clean schema edge cases +# --------------------------------------------------------------------------- + + +def test_clean_schema_non_dict(): + assert _clean_schema("string") == "string" + assert _clean_schema(42) == 42 + assert _clean_schema([1, {"$schema": "x"}, 3]) == [1, {}, 3] diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000..c8eb60c --- /dev/null +++ b/tools/__init__.py @@ -0,0 +1,2 @@ +from .registry import MAX_TOOL_ROUNDS, ToolCallRecord, get_tool_declarations # noqa: F401 +from .handlers import DEFAULT_TOOLS, get_registry # noqa: F401 diff --git a/tools/handlers.py b/tools/handlers.py new file mode 100644 index 0000000..c90204b --- /dev/null +++ b/tools/handlers.py @@ -0,0 +1,226 @@ +"""Tool handler implementations for the Gemini function-calling framework. + +Every handler is a sync callable that returns a dict (→ function_response). +All handlers are read-only against external systems. +""" + +import json +import logging +import os +import urllib.error +import urllib.parse +import urllib.request +from decimal import Decimal +from typing import Optional + +from pydantic import BaseModel, Field +from stellar_sdk import Server +from stellar_sdk.exceptions import NotFoundError +from stellar_sdk.strkey import StrKey + +from .registry import Tool, ToolRegistry + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Shared Stellar config (mirrors stellar.py) +# --------------------------------------------------------------------------- + +STELLAR_NETWORK = os.getenv("STELLAR_NETWORK", "testnet") + +HORIZON_URLS = { + "testnet": "https://horizon-testnet.stellar.org", + "public": "https://horizon.stellar.org", +} + +USDC_ISSUERS = { + "public": "GA5ZSEJYB37JRC5AVCIA5MOP4RHTM335X2KGX3IHOJAPP5RE34K4KZVN", + "testnet": "GBBD47IF6LWK7P7MDEVSCWR7DPUWV3NY3DTQEVFL4NAT4AQH3ZLLFLA5", +} + +ZAKAT_RATE = Decimal("0.025") +DEFAULT_NISAB_USD = Decimal(os.getenv("ZAKAT_NISAB_USD", "6000")) + +DISCLAIMER = ( + "This is an automated estimate based on your on-chain USDC balance only. " + "Zakat rulings depend on your full wealth, debts, and the hawl (lunar year). " + "Please consult a qualified scholar for a definitive ruling." +) + + +def _horizon_url() -> str: + return HORIZON_URLS.get(STELLAR_NETWORK, HORIZON_URLS["testnet"]) + + +def _usdc_issuer() -> str: + return USDC_ISSUERS.get(STELLAR_NETWORK, USDC_ISSUERS["testnet"]) + + +def _fetch_usdc_balance(public_key: str) -> Optional[Decimal]: + """Return the account's USDC balance, or None if it has no USDC trustline.""" + server = Server(_horizon_url()) + try: + account = server.accounts().account_id(public_key).call() + except NotFoundError: + return None + for balance in account.get("balances", []): + if ( + balance.get("asset_code") == "USDC" + and balance.get("asset_issuer") == _usdc_issuer() + ): + return Decimal(balance["balance"]) + return None + + +# --------------------------------------------------------------------------- +# Tool: calculate_zakat +# --------------------------------------------------------------------------- + + +class CalculateZakatArgs(BaseModel): + public_key: str = Field(..., description="Stellar account public key (G...)") + nisab_usd: Optional[float] = Field(None, description="Override nisab threshold in USD") + + +def _calculate_zakat(public_key: str, nisab_usd: Optional[float] = None) -> dict: + if not StrKey.is_valid_ed25519_public_key(public_key): + return {"error": "Invalid Stellar public key. Expected a 56-character key starting with G."} + + logger.info("Zakat lookup for %s on %s", public_key[:8], STELLAR_NETWORK) + balance = _fetch_usdc_balance(public_key) + nisab = Decimal(str(nisab_usd)) if nisab_usd else DEFAULT_NISAB_USD + + if balance is None: + return { + "network": STELLAR_NETWORK, + "public_key": public_key, + "has_usdc_trustline": False, + "usdc_balance": "0", + "nisab_usd": str(nisab), + "zakat_rate": str(ZAKAT_RATE), + "zakat_due": "0", + "message": "This account has no USDC trustline.", + "disclaimer": DISCLAIMER, + } + + meets_nisab = balance >= nisab + zakat_due = (balance * ZAKAT_RATE).quantize(Decimal("0.0000001")) if meets_nisab else Decimal("0") + + if meets_nisab: + message = ( + f"Your USDC balance of {balance} meets the nisab threshold of {nisab} USD. " + f"If held for a full lunar year, the zakat due is {zakat_due} USDC (2.5%)." + ) + else: + message = ( + f"Your USDC balance of {balance} is below the nisab threshold of {nisab} USD, " + "so no zakat is due on this balance alone." + ) + + return { + "network": STELLAR_NETWORK, + "public_key": public_key, + "has_usdc_trustline": True, + "usdc_balance": str(balance), + "nisab_usd": str(nisab), + "zakat_rate": str(ZAKAT_RATE), + "zakat_due": str(zakat_due), + "message": message, + "disclaimer": DISCLAIMER, + } + + +CALCULATE_ZAKAT_TOOL = Tool( + name="calculate_zakat", + description=( + "Calculate the zakat (obligatory charity) due on a Stellar wallet's " + "on-chain USDC balance. Returns the balance, nisab threshold, " + "zakat rate (2.5%), and the calculated amount due. Includes a " + "scholar-consultation disclaimer." + ), + args_schema=CalculateZakatArgs, + handler=_calculate_zakat, + timeout_seconds=15, +) + +# --------------------------------------------------------------------------- +# Tool: get_stellar_info +# --------------------------------------------------------------------------- + + +class GetStellarInfoArgs(BaseModel): + pass + + +def _get_stellar_info() -> dict: + return { + "network": STELLAR_NETWORK, + "horizon": _horizon_url(), + "usdc_issuer": _usdc_issuer(), + "features": ["zakat"], + } + + +GET_STELLAR_INFO_TOOL = Tool( + name="get_stellar_info", + description="Get configuration information about the Stellar network integration used by Deen Bridge.", + args_schema=GetStellarInfoArgs, + handler=_get_stellar_info, + timeout_seconds=5, +) + +# --------------------------------------------------------------------------- +# Tool: search_courses +# --------------------------------------------------------------------------- + + +class SearchCoursesArgs(BaseModel): + query: str = Field(..., description="Search query for course titles or descriptions") + + +BACKEND_API_URL = os.getenv("BACKEND_API_URL", "https://dnb-backend-api.onrender.com") + + +def _search_courses(query: str) -> dict: + url = f"{BACKEND_API_URL.rstrip('/')}/api/courses/search?q={urllib.parse.quote(query)}" + logger.info("Course search query=%s url=%s", query[:80], url) + try: + req = urllib.request.Request(url, method="GET") + with urllib.request.urlopen(req, timeout=8) as resp: + data = json.loads(resp.read().decode()) + return {"results": data, "count": len(data) if isinstance(data, list) else 1} + except urllib.error.HTTPError as e: + return {"error": f"Backend returned HTTP {e.code}: {e.reason}"} + except urllib.error.URLError as e: + return {"error": f"Cannot reach course backend: {e.reason}"} + except Exception as e: + return {"error": f"Course search failed: {e}"} + + +SEARCH_COURSES_TOOL = Tool( + name="search_courses", + description=( + "Search for Islamic courses available on the Deen Bridge platform. " + "Returns matching course titles, descriptions, and enrollment details." + ), + args_schema=SearchCoursesArgs, + handler=_search_courses, + timeout_seconds=10, +) + +# --------------------------------------------------------------------------- +# Default registry +# --------------------------------------------------------------------------- + +DEFAULT_TOOLS = [ + CALCULATE_ZAKAT_TOOL, + GET_STELLAR_INFO_TOOL, + SEARCH_COURSES_TOOL, +] + + +def get_registry() -> ToolRegistry: + reg = ToolRegistry() + for tool in DEFAULT_TOOLS: + reg.register(tool) + return reg diff --git a/tools/registry.py b/tools/registry.py new file mode 100644 index 0000000..9cc061d --- /dev/null +++ b/tools/registry.py @@ -0,0 +1,139 @@ +"""Tool definition, registry, and schema bridging for Gemini function-calling. + +Safety +------ +Tool handlers must be read-only against external systems. +Nothing outside the explicit allowlist is ever callable by the model. +Every invocation is logged with name, truncated args, duration, and outcome. +""" + +import logging +import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + +MAX_TOOL_ROUNDS = 4 +_EXECUTOR = ThreadPoolExecutor(max_workers=4) + + +# --------------------------------------------------------------------------- +# Data types +# --------------------------------------------------------------------------- + + +@dataclass +class Tool: + name: str + description: str + args_schema: type[BaseModel] + handler: Callable[..., dict[str, Any]] + timeout_seconds: int = 10 + + +class ToolCallRecord(BaseModel): + tool_name: str + args: dict[str, Any] + result: str + + +# --------------------------------------------------------------------------- +# Schema bridging — pydantic v2 model_json_schema() -> Gemini declaration +# --------------------------------------------------------------------------- + +_UNSUPPORTED_KEYS = frozenset({"$schema", "$defs", "title"}) + + +def _clean_schema(d: Any) -> Any: + """Recursively strip keys Gemini's FunctionDeclaration does not accept.""" + if isinstance(d, dict): + return {k: _clean_schema(v) for k, v in d.items() if k not in _UNSUPPORTED_KEYS} + if isinstance(d, list): + return [_clean_schema(v) for v in d] + return d + + +def get_tool_declarations(tools: list[Tool]) -> list[dict[str, Any]]: + """Build Gemini FunctionDeclaration dicts from a list of Tool objects.""" + declarations = [] + for tool in tools: + raw = tool.args_schema.model_json_schema() + cleaned = _clean_schema(raw) + declarations.append({ + "name": tool.name, + "description": tool.description, + "parameters": cleaned, + }) + return declarations + + +# --------------------------------------------------------------------------- +# Tool registry +# --------------------------------------------------------------------------- + + +class ToolRegistry: + """Explicit allowlist of tools callable by the model.""" + + def __init__(self) -> None: + self._tools: dict[str, Tool] = {} + + def register(self, tool: Tool) -> None: + if tool.name in self._tools: + raise ValueError(f"Tool '{tool.name}' is already registered") + self._tools[tool.name] = tool + + def get(self, name: str) -> Optional[Tool]: + return self._tools.get(name) + + @property + def all(self) -> list[Tool]: + return list(self._tools.values()) + + def declarations(self) -> list[dict[str, Any]]: + return get_tool_declarations(self.all) + + +# --------------------------------------------------------------------------- +# Tool execution helpers +# --------------------------------------------------------------------------- + + +def run_tool_handler(tool: Tool, args: dict[str, Any]) -> dict[str, Any]: + """Validate args through pydantic and execute the handler with timeout.""" + try: + validated = tool.args_schema.model_validate(args) + validated_args = validated.model_dump(exclude_unset=True) + except Exception as exc: + return {"error": f"Argument validation failed for '{tool.name}': {exc}"} + + started = time.perf_counter() + logger.info("Tool call: %s args=%s", tool.name, _truncate(validated_args, 200)) + + try: + future = _EXECUTOR.submit(tool.handler, **validated_args) + result = future.result(timeout=tool.timeout_seconds) + except FutureTimeout: + result = {"error": f"Tool '{tool.name}' timed out after {tool.timeout_seconds}s"} + logger.warning("Tool timeout: %s (%ss)", tool.name, tool.timeout_seconds) + except Exception as exc: + result = {"error": f"Tool '{tool.name}' error: {exc}"} + logger.error("Tool error: %s — %s", tool.name, exc) + + duration = time.perf_counter() - started + logger.info( + "Tool result: %s duration=%.2fs outcome=%s", + tool.name, + duration, + "error" if "error" in result else "ok", + ) + return result + + +def _truncate(obj: Any, limit: int = 200) -> Any: + s = str(obj) + return s[:limit] + "..." if len(s) > limit else s