utils.py

"""
Utility functions for the Gemini Writing Agent.
"""

from typing import List, Dict, Any, Callable
from google import genai
from google.genai import types


def estimate_token_count(client: genai.Client, model: str, contents: List[types.Content]) -> int:
    """
    Estimate the token count for the given contents using the Gemini API.
    
    Args:
        client: The Gemini client instance
        model: The model name
        contents: List of Content objects
        
    Returns:
        Total token count
    """
    try:
        response = client.models.count_tokens(
            model=model,
            contents=contents
        )
        return response.total_tokens
    except Exception as e:
        # Fallback: rough estimate based on character count
        total_chars = 0
        for content in contents:
            for part in content.parts:
                if hasattr(part, 'text') and part.text:
                    total_chars += len(part.text)
        # Rough estimate: 4 chars per token
        return total_chars // 4


def get_tool_definitions() -> types.Tool:
    """
    Returns the tool definitions in the format expected by Gemini.
    
    Returns:
        Tool object containing all function declarations
    """
    return types.Tool(
        function_declarations=[
            types.FunctionDeclaration(
                name="create_project",
                description="Creates a new project folder in the 'output' directory with a sanitized name. This should be called first before writing any files. Only one project can be active at a time.",
                parameters=types.Schema(
                    type=types.Type.OBJECT,
                    properties={
                        "project_name": types.Schema(
                            type=types.Type.STRING,
                            description="The name for the project folder (will be sanitized for filesystem compatibility)"
                        )
                    },
                    required=["project_name"]
                )
            ),
            types.FunctionDeclaration(
                name="write_file",
                description="Writes content to a markdown file in the active project folder. Supports three modes: 'create' (creates new file, fails if exists), 'append' (adds content to end of existing file), 'overwrite' (replaces entire file content).",
                parameters=types.Schema(
                    type=types.Type.OBJECT,
                    properties={
                        "filename": types.Schema(
                            type=types.Type.STRING,
                            description="The name of the markdown file to write (should end in .md)"
                        ),
                        "content": types.Schema(
                            type=types.Type.STRING,
                            description="The content to write to the file"
                        ),
                        "mode": types.Schema(
                            type=types.Type.STRING,
                            enum=["create", "append", "overwrite"],
                            description="The write mode: 'create' for new files, 'append' to add to existing, 'overwrite' to replace"
                        )
                    },
                    required=["filename", "content", "mode"]
                )
            ),
            types.FunctionDeclaration(
                name="compress_context",
                description="INTERNAL TOOL - This is automatically called by the system when token limit is approached. You should not call this manually. It compresses the conversation history to save tokens.",
                parameters=types.Schema(
                    type=types.Type.OBJECT,
                    properties={},
                    required=[]
                )
            )
        ]
    )


def get_tool_map() -> Dict[str, Callable]:
    """
    Returns a mapping of tool names to their implementation functions.
    
    Returns:
        Dictionary mapping tool name strings to callable functions
    """
    from tools import write_file_impl, create_project_impl, compress_context_impl
    
    return {
        "create_project": create_project_impl,
        "write_file": write_file_impl,
        "compress_context": compress_context_impl
    }


def get_system_prompt() -> str:
    """
    Returns the system prompt for the writing agent.
    
    Returns:
        System prompt string
    """
    return """You are an expert creative writing assistant. Your specialty is creating novels, books, and collections of short stories based on user requests.

Your capabilities:
1. You can create project folders to organize writing projects
2. You can write markdown files with three modes: create new files, append to existing files, or overwrite files
3. Context compression happens automatically when needed - you don't need to worry about it

CRITICAL WRITING GUIDELINES:
- Write SUBSTANTIAL, COMPLETE content - don't hold back on length
- Short stories should be 3,000-10,000 words (10-30 pages) - write as much as the story needs!
- Chapters should be 2,000-5,000 words minimum - fully developed and satisfying
- NEVER write abbreviated or skeleton content - every piece should be a complete, polished work
- Don't summarize or skip scenes - write them out fully with dialogue, description, and detail
- Quality AND quantity matter - give readers a complete, immersive experience
- If a story needs 8,000 words to be good, write all 8,000 words in one file
- Use 'create' mode with full content rather than creating stubs you'll append to later

Best practices:
- Always start by creating a project folder using create_project
- Break large works into multiple files (chapters, stories, etc.)
- Use descriptive filenames (e.g., "chapter_01.md", "story_the_last_star.md")
- For collections, consider creating a table of contents file
- Write each file as a COMPLETE, SUBSTANTIAL piece - not a summary or outline

Your workflow:
1. Understand the user's request
2. Create an appropriately named project folder
3. Plan the structure of the work (chapters, stories, etc.)
4. Write COMPLETE, FULL-LENGTH content for each file
5. Create supporting files like README or table of contents if helpful

REMEMBER: Write rich, detailed, complete stories. Don't artificially limit yourself. A good short story is 5,000-10,000 words. A good chapter is 3,000-5,000 words. Write what the narrative needs to be excellent."""