google-genai[patch]: fix tool format, use protos (#17284)

This commit is contained in:
Erick Friis 2024-02-08 19:36:49 -08:00 committed by GitHub
parent d8913b9428
commit febf9540b9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -7,6 +7,7 @@ from typing import (
Union, Union,
) )
import google.ai.generativelanguage as glm
from langchain_core.pydantic_v1 import BaseModel from langchain_core.pydantic_v1 import BaseModel
from langchain_core.tools import BaseTool from langchain_core.tools import BaseTool
from langchain_core.utils.json_schema import dereference_refs from langchain_core.utils.json_schema import dereference_refs
@ -14,55 +15,36 @@ from langchain_core.utils.json_schema import dereference_refs
FunctionCallType = Union[BaseTool, Type[BaseModel], Dict] FunctionCallType = Union[BaseTool, Type[BaseModel], Dict]
TYPE_ENUM = { TYPE_ENUM = {
"string": 1, "string": glm.Type.STRING,
"number": 2, "number": glm.Type.NUMBER,
"integer": 3, "integer": glm.Type.INTEGER,
"boolean": 4, "boolean": glm.Type.BOOLEAN,
"array": 5, "array": glm.Type.ARRAY,
"object": 6, "object": glm.Type.OBJECT,
} }
def convert_to_genai_function_declarations( def convert_to_genai_function_declarations(
function_calls: List[FunctionCallType], function_calls: List[FunctionCallType],
) -> Dict: ) -> List[glm.Tool]:
function_declarations = [] return [
for fc in function_calls: glm.Tool(
function_declarations.append(_convert_to_genai_function(fc)) function_declarations=[_convert_to_genai_function(fc)],
return { )
"function_declarations": function_declarations, for fc in function_calls
} ]
def _convert_to_genai_function(fc: FunctionCallType) -> Dict: def _convert_to_genai_function(fc: FunctionCallType) -> glm.FunctionDeclaration:
"""
Produce
{
"name": "get_weather",
"description": "Determine weather in my location",
"parameters": {
"properties": {
"location": {
"description": "The city and state e.g. San Francisco, CA",
"type_": 1
},
"unit": { "enum": ["c", "f"], "type_": 1 }
},
"required": ["location"],
"type_": 6
}
}
"""
if isinstance(fc, BaseTool): if isinstance(fc, BaseTool):
return _convert_tool_to_genai_function(fc) return _convert_tool_to_genai_function(fc)
elif isinstance(fc, type) and issubclass(fc, BaseModel): elif isinstance(fc, type) and issubclass(fc, BaseModel):
return _convert_pydantic_to_genai_function(fc) return _convert_pydantic_to_genai_function(fc)
elif isinstance(fc, dict): elif isinstance(fc, dict):
return { return glm.FunctionDeclaration(
**fc, name=fc["name"],
"parameters": { description=fc.get("description"),
parameters={
"properties": { "properties": {
k: { k: {
"type_": TYPE_ENUM[v["type"]], "type_": TYPE_ENUM[v["type"]],
@ -73,20 +55,20 @@ def _convert_to_genai_function(fc: FunctionCallType) -> Dict:
"required": fc["parameters"].get("required", []), "required": fc["parameters"].get("required", []),
"type_": TYPE_ENUM[fc["parameters"]["type"]], "type_": TYPE_ENUM[fc["parameters"]["type"]],
}, },
} )
else: else:
raise ValueError(f"Unsupported function call type {fc}") raise ValueError(f"Unsupported function call type {fc}")
def _convert_tool_to_genai_function(tool: BaseTool) -> Dict: def _convert_tool_to_genai_function(tool: BaseTool) -> glm.FunctionDeclaration:
if tool.args_schema: if tool.args_schema:
schema = dereference_refs(tool.args_schema.schema()) schema = dereference_refs(tool.args_schema.schema())
schema.pop("definitions", None) schema.pop("definitions", None)
return { return glm.FunctionDeclaration(
"name": tool.name or schema["title"], name=tool.name or schema["title"],
"description": tool.description or schema["description"], description=tool.description or schema["description"],
"parameters": { parameters={
"properties": { "properties": {
k: { k: {
"type_": TYPE_ENUM[v["type"]], "type_": TYPE_ENUM[v["type"]],
@ -97,31 +79,30 @@ def _convert_tool_to_genai_function(tool: BaseTool) -> Dict:
"required": schema["required"], "required": schema["required"],
"type_": TYPE_ENUM[schema["type"]], "type_": TYPE_ENUM[schema["type"]],
}, },
} )
else: else:
return { return glm.FunctionDeclaration(
"name": tool.name, name=tool.name,
"description": tool.description, description=tool.description,
"parameters": { parameters={
"properties": { "properties": {
"__arg1": {"type": "string"}, "__arg1": {"type_": TYPE_ENUM["string"]},
}, },
"required": ["__arg1"], "required": ["__arg1"],
"type_": TYPE_ENUM["object"], "type_": TYPE_ENUM["object"],
}, },
} )
def _convert_pydantic_to_genai_function( def _convert_pydantic_to_genai_function(
pydantic_model: Type[BaseModel], pydantic_model: Type[BaseModel],
) -> Dict: ) -> glm.FunctionDeclaration:
schema = dereference_refs(pydantic_model.schema()) schema = dereference_refs(pydantic_model.schema())
schema.pop("definitions", None) schema.pop("definitions", None)
return glm.FunctionDeclaration(
return { name=schema["title"],
"name": schema["title"], description=schema.get("description", ""),
"description": schema.get("description", ""), parameters={
"parameters": {
"properties": { "properties": {
k: { k: {
"type_": TYPE_ENUM[v["type"]], "type_": TYPE_ENUM[v["type"]],
@ -132,4 +113,4 @@ def _convert_pydantic_to_genai_function(
"required": schema["required"], "required": schema["required"],
"type_": TYPE_ENUM[schema["type"]], "type_": TYPE_ENUM[schema["type"]],
}, },
} )