diff --git a/libs/vertexai/langchain_google_vertexai/utils.py b/libs/vertexai/langchain_google_vertexai/utils.py index 5b31ff0f..c9286a04 100644 --- a/libs/vertexai/langchain_google_vertexai/utils.py +++ b/libs/vertexai/langchain_google_vertexai/utils.py @@ -60,7 +60,7 @@ def create_context_cache( tool_config = _format_tool_config(tool_config) if tools is not None: - tools = _format_to_gapic_tool(tools) + tools = [_format_to_gapic_tool(tools)] cached_content = caching.CachedContent.create( model_name=model.full_model_name, diff --git a/libs/vertexai/tests/integration_tests/test_chains.py b/libs/vertexai/tests/integration_tests/test_chains.py index 295219f2..5f79a3c0 100644 --- a/libs/vertexai/tests/integration_tests/test_chains.py +++ b/libs/vertexai/tests/integration_tests/test_chains.py @@ -44,7 +44,7 @@ def test_create_structured_runnable() -> None: @pytest.mark.release def test_create_structured_runnable_with_prompt() -> None: - llm = ChatVertexAI(model_name=_DEFAULT_MODEL_NAME) + llm = ChatVertexAI(model_name=_DEFAULT_MODEL_NAME, temperature=0) prompt = ChatPromptTemplate.from_template( "Describe a random {class} and mention their name, {attr} and favorite food" ) diff --git a/libs/vertexai/tests/integration_tests/test_chat_models.py b/libs/vertexai/tests/integration_tests/test_chat_models.py index cca07998..e610f928 100644 --- a/libs/vertexai/tests/integration_tests/test_chat_models.py +++ b/libs/vertexai/tests/integration_tests/test_chat_models.py @@ -902,6 +902,76 @@ def test_context_catching(): assert isinstance(response.content, str) +@pytest.mark.extended +def test_context_catching_tools(): + from langchain import agents + + @tool + def get_secret_number() -> int: + """Gets secret number.""" + return 747 + + tools = [get_secret_number] + system_instruction = """ + You are an expert researcher. You always stick to the facts in the sources + provided, and never make up new facts. + + You have a get_secret_number function available. Use this tool if someone asks + for the secret number. + + Now look at these research papers, and answer the following questions. + + """ + + cached_content = create_context_cache( + model=ChatVertexAI( + model_name="gemini-1.5-pro-001", + ), + messages=[ + SystemMessage(content=system_instruction), + HumanMessage( + content=[ + { + "type": "image_url", + "image_url": { + "url": "gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf", + }, + }, + { + "type": "image_url", + "image_url": { + "url": "gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf" + }, + }, + ] + ), + ], + tools=tools, + ) + + chat = ChatVertexAI( + model_name="gemini-1.5-pro-001", + cached_content=cached_content, + ) + + prompt = ChatPromptTemplate.from_messages( + [ + ("human", "{input}"), + ("placeholder", "{agent_scratchpad}"), + ] + ) + agent = agents.create_tool_calling_agent( + llm=chat, + tools=tools, + prompt=prompt, + ) + agent_executor = agents.AgentExecutor( # type: ignore[call-arg] + agent=agent, tools=tools, verbose=False, stream_runnable=False + ) + response = agent_executor.invoke({"input": "what is the secret number?"}) + assert isinstance(response["output"], str) + + @pytest.mark.release def test_json_serializable() -> None: llm = ChatVertexAI( diff --git a/libs/vertexai/tests/integration_tests/test_standard.py b/libs/vertexai/tests/integration_tests/test_standard.py index bf54ba89..a6dfc60a 100644 --- a/libs/vertexai/tests/integration_tests/test_standard.py +++ b/libs/vertexai/tests/integration_tests/test_standard.py @@ -37,7 +37,11 @@ def chat_model_class(self) -> Type[BaseChatModel]: @property def chat_model_params(self) -> dict: - return {"model_name": "gemini-1.5-pro-001", "rate_limiter": rate_limiter} + return { + "model_name": "gemini-1.5-pro-001", + "rate_limiter": rate_limiter, + "temperature": 0, + } @property def supports_image_inputs(self) -> bool: