test_stream_tool_gemini.py

import os,json
from litellm import completion

os.environ["GEMINI_API_KEY"] = "gemini-key"
os.environ["ANTHROPIC_API_KEY"] = "anthropic-key"



SYSTEM_PROMPT = "You are an AI assistant"

tools = [
    { 
        "type": "function",
        "function": {
            "name": "generate_series_of_questions",
            "description": "Generate a series of questions, given a topic.",
            "parameters": {
                "type": "object",
                "properties": {
                    "questions": {
                        "type": "array",
                        "description": "The questions to be generated.",                                
                        "items" : {}
                    },
                },
                "required": ["questions"]
            },
        }
    },
]


messages = [
    {'role': 'system', 'content': SYSTEM_PROMPT},
    {'role': 'user', 'content': "Propose a general topic and then generate three questions it. When generating the questions, call tje function generate_series_of_questions."},
]


response =  completion(
                    model="gemini/gemini-1.5-pro",
                    # model="claude-3-5-sonnet-20240620",
                    messages=messages,
                    stream=True,
                    temperature=0.75,
                    tools=tools,
                    stream_options = {"include_usage": True},
                    )

tool_str = ""
buff_mem = []
text = []
for tk in response:
    buff_mem.append(tk)
    delta = tk.choices[0].delta if tk.choices else None 
    if delta and delta.content:
        text.append(delta.content)
    elif delta and delta.tool_calls:
        tool_str += delta.tool_calls[0].function.arguments
print("".join(text))
if tool_str: # call the tool
    json_tool = json.loads(tool_str)
    questions = json_tool['questions']
    print("---")
    print("questions:",questions)