import genkit as gk
from genkit.config import GeminiProLLMConfig
from a2a.server import A2AServer
 
# Initialize the Vertex GenKit model
llm_config = GeminiProLLMConfig(
    safety_settings={
        gk.SafetyCategory.HARASSMENT: gk.SafetyThreshold.BLOCK_MEDIUM_AND_ABOVE,
        gk.SafetyCategory.HATE_SPEECH: gk.SafetyThreshold.BLOCK_MEDIUM_AND_ABOVE,
    }
)
model = gk.llm(llm_config)
 
# Define a function to process user requests with GenKit
def process_with_genkit(user_input: str) -> str:
    # Create a conversation
    conversation = model.start_chat()
    
    # Send the user's input to the model
    response = conversation.send_message(user_input)
    
    # Return the model's response
    return response.text
 
# Create an A2A server
a2a_server = A2AServer(
    agent=process_with_genkit,
    agent_card={
        "name": "GenKit Assistant",
        "description": "A helpful assistant powered by Google's Vertex GenKit",
        "skills": [
            {
                "id": "general-assistance",
                "name": "General Assistance",
                "description": "Can provide information and assistance on a wide range of topics"
            }
        ]
    }
)
 
# Run the server
if __name__ == "__main__":
    a2a_server.run(host="0.0.0.0", port=8000)

import genkit as gk
from genkit.config import GeminiProLLMConfig
from a2a.server import A2AServer
import json
 
# Define a simple weather tool
def get_weather(location: str) -> str:
    return f"The weather in {location} is sunny with a high of 75°F."
 
# Initialize the model with tool support
llm_config = GeminiProLLMConfig()
model = gk.llm(llm_config)
 
# Define the tool schema
weather_tool = gk.Tool(
    name="get_weather",
    description="Get the current weather for a location",
    function=get_weather,
    parameters={
        "type": "object",
        "properties": {
            "location": {
                "type": "string",
                "description": "The city and state/country"
            }
        },
        "required": ["location"]
    }
)
 
# Function to process requests with tool support
def process_with_tools(user_input: str) -> str:
    # Create a conversation with tools
    conversation = model.start_chat(tools=[weather_tool])
    
    # Send the user message
    response = conversation.send_message(user_input)
    
    # Handle any tool calls
    while response.tool_calls:
        for tool_call in response.tool_calls:
            # Call the tool and get its result
            if tool_call.name == "get_weather":
                args = json.loads(tool_call.args)
                result = get_weather(args.get("location", "unknown"))
                # Send the tool result back to the model
                response = conversation.send_tool_result(
                    tool_call.id, 
                    result
                )
    
    # Return the final text response
    return response.text
 
# Create the A2A server with tool support
a2a_server = A2AServer(
    agent=process_with_tools,
    agent_card={
        "name": "GenKit Weather Assistant",
        "description": "An assistant that can provide weather information using Google's Vertex GenKit",
        "skills": [
            {
                "id": "weather-lookup",
                "name": "Weather Information",
                "description": "Can provide current weather conditions for any location"
            }
        ]
    }
)
 
# Run the server
if __name__ == "__main__":
    a2a_server.run(host="0.0.0.0", port=8000)