diff --git a/examples/langchain/chat/chat_api_invoke.py b/examples/langchain/chat/chat_api_invoke.py new file mode 100644 index 00000000..a266c4b7 --- /dev/null +++ b/examples/langchain/chat/chat_api_invoke.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +"""LangChain invoke example using Chat API.""" + +from __future__ import annotations + +from langchain_core.messages import AIMessage, HumanMessage +from yandex_ai_studio_sdk import AIStudio + + +def main() -> None: + # You can set authentication using environment variables instead of the 'auth' argument: + # YC_OAUTH_TOKEN, YC_TOKEN, YC_IAM_TOKEN, or YC_API_KEY + # You can also set 'folder_id' using the YC_FOLDER_ID environment variable + sdk = AIStudio( + # folder_id="", + # auth="", + ) + sdk.setup_default_logging() + + model = sdk.chat.completions('yandexgpt').langchain(timeout=60) + + result = model.invoke( + [ + HumanMessage(content="hello!"), + AIMessage(content="Hi there human!"), + HumanMessage(content="Meow!"), + ] + ) + print(result) + print(f"Content: {result.content}") + print(f"Usage: {result.usage_metadata}") + print(f"Finish reason: {result.response_metadata.get('finish_reason')}") + + +if __name__ == '__main__': + main() diff --git a/src/yandex_ai_studio_sdk/_chat/completions/langchain.py b/src/yandex_ai_studio_sdk/_chat/completions/langchain.py new file mode 100644 index 00000000..e720955d --- /dev/null +++ b/src/yandex_ai_studio_sdk/_chat/completions/langchain.py @@ -0,0 +1,252 @@ +"""LangChain integration for Yandex AI Studio Chat API. + +This module is optional: requires ``langchain_core`` to be installed. +Provides :class:`ChatYandexGPT` — a LangChain ``BaseChatModel`` adapter +that uses the Chat API backend. +""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator, Iterator +from typing import Any + +from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun +from langchain_core.language_models.chat_models import BaseChatModel as LCBaseChatModel +from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.messages.ai import UsageMetadata +from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult +from yandex_ai_studio_sdk._types.langchain import BaseYandexLanguageModel +from yandex_ai_studio_sdk._utils.langchain import make_async_run_manager +from yandex_ai_studio_sdk._utils.sync import run_sync_generator_impl, run_sync_impl + +from .model import BaseChatModel as ChatAPIModel +from .result import ChatChoice, ChatModelResult, DeltaChatChoice + +# ========================================================================= +# Message conversion: LangChain → Chat API +# ========================================================================= + + +def _transform_messages(messages: list[BaseMessage]) -> list[dict[str, Any]]: + """Convert LangChain messages to Chat API dict format.""" + result: list[dict[str, Any]] = [] + + for message in messages: + if isinstance(message, ToolMessage): + result.append({ + "role": "tool", + "content": str(message.content), + "tool_call_id": message.tool_call_id, + }) + + elif isinstance(message, AIMessage) and message.tool_calls: + result.append({ + "role": "assistant", + "tool_calls": [ + { + "id": tc["id"], + "type": "function", + "function": { + "name": tc["name"], + "arguments": json.dumps( + tc["args"], ensure_ascii=False, + ), + }, + } + for tc in message.tool_calls + ], + }) + + elif isinstance(message, HumanMessage): + content = message.content + if isinstance(content, list): + # Multimodal: list of content parts + result.append({"role": "user", "content": content}) + else: + result.append({"role": "user", "content": str(content)}) + + elif isinstance(message, AIMessage): + result.append({"role": "assistant", "content": str(message.content)}) + + elif isinstance(message, SystemMessage): + result.append({"role": "system", "content": str(message.content)}) + + return result + + +# ========================================================================= +# Response parsing: Chat API → LangChain +# ========================================================================= + + +def _parse_tool_calls(choice: ChatChoice) -> list[dict[str, Any]]: + """Extract tool calls from ChatChoice into LangChain format.""" + if not choice.tool_calls: + return [] + + return [ + { + "id": tc.id or "", + "name": tc.function.name, + "args": tc.function.arguments, + "type": "tool_call", + } + for tc in choice.tool_calls + if tc.function # skip malformed tool calls + ] + + +def _make_usage(result: ChatModelResult) -> UsageMetadata | None: + """Build LangChain UsageMetadata from Chat API usage stats.""" + if not result.usage: + return None + return UsageMetadata( + input_tokens=result.usage.input_text_tokens, + output_tokens=result.usage.completion_tokens, + total_tokens=result.usage.total_tokens, + ) + + +# ========================================================================= +# ChatYandexGPT — LangChain adapter for Chat API +# ========================================================================= + + +class ChatYandexGPT(BaseYandexLanguageModel[ChatAPIModel], LCBaseChatModel): + """LangChain chat model for Yandex GPT via Chat API. + + Supports text, tool calls, tool results, multimodal content, and streaming. + + Example: + >>> sdk = AIStudio() + >>> model = sdk.chat.completions('yandexgpt').langchain() + >>> result = model.invoke([HumanMessage("Hello!")]) + """ + + class Config: + arbitrary_types_allowed = True + + @property + def _sdk(self): + return self.ycmlsdk_model._sdk + + # ----------------------------------------------------------------- + # Sync → async delegation (same pattern as legacy ChatYandexGPT) + # ----------------------------------------------------------------- + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + async_rm = make_async_run_manager(run_manager) if run_manager else None + return run_sync_impl( + self._agenerate(messages, stop, async_rm, **kwargs), + self._sdk, + ) + + def _stream( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> Iterator[ChatGenerationChunk]: + async_rm = make_async_run_manager(run_manager) if run_manager else None + return run_sync_generator_impl( + self._astream(messages, stop, async_rm, **kwargs), + self._sdk, + ) + + # ----------------------------------------------------------------- + # Core: invoke + # ----------------------------------------------------------------- + + async def _agenerate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + chat_messages = _transform_messages(messages) + + sdk_result: ChatModelResult = await self.ycmlsdk_model._run( + messages=chat_messages, + timeout=self.timeout, + ) + + usage = _make_usage(sdk_result) + + generations: list[ChatGeneration] = [] + for choice in sdk_result.choices: + tool_calls = _parse_tool_calls(choice) + + ai_message = AIMessage( + content=choice.text or "", + tool_calls=tool_calls or [], + usage_metadata=usage, + response_metadata={ + "finish_reason": choice.finish_reason.value, + "model": sdk_result.model, + "status": choice.status.name, + }, + ) + generations.append(ChatGeneration(message=ai_message)) + + return ChatResult( + generations=generations, + llm_output={"model": sdk_result.model, "id": sdk_result.id}, + ) + + # ----------------------------------------------------------------- + # Core: stream + # ----------------------------------------------------------------- + + async def _astream( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> AsyncIterator[ChatGenerationChunk]: + chat_messages = _transform_messages(messages) + + async for sdk_result in self.ycmlsdk_model._run_stream( + messages=chat_messages, + timeout=self.timeout, + ): + choice = sdk_result.choices[0] + usage = _make_usage(sdk_result) + + delta = choice.delta if isinstance(choice, DeltaChatChoice) else choice.text + tool_calls = _parse_tool_calls(choice) + + # tool_call_chunks: args as JSON string (LangChain convention) + tool_call_chunks = [ + { + "id": tc["id"], + "name": tc["name"], + "args": json.dumps(tc["args"], ensure_ascii=False), + "index": i, + "type": "tool_call_chunk", + } + for i, tc in enumerate(tool_calls) + ] if tool_calls else [] + + chunk = AIMessageChunk( + content=delta or "", + tool_call_chunks=tool_call_chunks, + usage_metadata=usage, + response_metadata={ + "finish_reason": choice.finish_reason.value, + "status": choice.status.name, + }, + ) + yield ChatGenerationChunk(message=chunk) + + +ChatYandexGPT.model_rebuild() diff --git a/src/yandex_ai_studio_sdk/_chat/completions/model.py b/src/yandex_ai_studio_sdk/_chat/completions/model.py index 2e9d40ba..e1940dc4 100644 --- a/src/yandex_ai_studio_sdk/_chat/completions/model.py +++ b/src/yandex_ai_studio_sdk/_chat/completions/model.py @@ -79,6 +79,17 @@ def configure( # type: ignore[override] extra_query=extra_query, ) + def langchain(self, timeout: int = 60): + """Initialize a LangChain chat model adapter. + + Requires ``langchain_core`` to be installed. + + :param timeout: Default timeout in seconds. Defaults to 60. + """ + from .langchain import ChatYandexGPT # pylint: disable=import-outside-toplevel + + return ChatYandexGPT(ycmlsdk_model=self, timeout=timeout) + def _build_request_json(self, messages: ChatMessageInputType, stream: bool) -> dict[str, Any]: result = { 'model': self._uri, diff --git a/tests/langchain_/cassettes/test_chat_api/test_ainvoke.yaml b/tests/langchain_/cassettes/test_chat_api/test_ainvoke.yaml new file mode 100644 index 00000000..59fef63a --- /dev/null +++ b/tests/langchain_/cassettes/test_chat_api/test_ainvoke.yaml @@ -0,0 +1,49 @@ +interactions: +- request: + body: '{"model":"gpt://b1gr2p0etomug2vas4ie/yandexgpt/latest","messages":[{"content":"hello!","role":"user"},{"content":"Hi + there human!","role":"assistant"},{"content":"Meow!","role":"user"}],"stream":false}' + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate, zstd + Connection: + - keep-alive + Content-Length: + - '201' + Content-Type: + - application/json + Host: + - llm.api.cloud.yandex.net + User-Agent: + - yandex-ai-studio-sdk/0.20.0 python/3.13 + x-client-request-id: + - 1b3c1788-a8b0-499f-92bb-3fe7523a4389 + method: POST + uri: https://llm.api.cloud.yandex.net/v1/chat/completions + response: + body: + string: '{"id":"f2bb6ba4-8f88-4ac9-87c0-6b41db01c635","object":"chat.completion","created":1774968384,"model":"gpt://b1gr2p0etomug2vas4ie/yandexgpt/latest","choices":[{"index":0,"message":{"role":"assistant","content":"Are + you trying to tell me that you''re a cat? Or are you just in the mood for + some feline fun? In any case, meow right back at you! Meow!"},"finish_reason":"stop"}],"usage":{"prompt_tokens":29,"total_tokens":69,"completion_tokens":40,"prompt_tokens_details":{"cached_tokens":0}}} + + ' + headers: + content-length: + - '489' + content-type: + - application/json + date: + - Tue, 31 Mar 2026 14:46:24 GMT + server: + - ycalb + x-client-request-id: + - 1b3c1788-a8b0-499f-92bb-3fe7523a4389 + x-request-id: + - 1ff57598-fabc-40bf-89ca-ace621e49c22 + x-server-trace-id: + - 85cfd10b60cdef2f:a32c621e24a87b57:55f05b16925bf389:1 + status: + code: 200 + message: OK +version: 1 diff --git a/tests/langchain_/cassettes/test_chat_api/test_astream.yaml b/tests/langchain_/cassettes/test_chat_api/test_astream.yaml new file mode 100644 index 00000000..9ef93a30 --- /dev/null +++ b/tests/langchain_/cassettes/test_chat_api/test_astream.yaml @@ -0,0 +1,61 @@ +interactions: +- request: + body: '{"model":"gpt://b1gr2p0etomug2vas4ie/yandexgpt/latest","messages":[{"content":"hello!","role":"user"},{"content":"Hi + there human!","role":"assistant"},{"content":"Meow!","role":"user"}],"stream":true}' + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate, zstd + Connection: + - keep-alive + Content-Length: + - '200' + Content-Type: + - application/json + Host: + - llm.api.cloud.yandex.net + User-Agent: + - yandex-ai-studio-sdk/0.20.0 python/3.13 + x-client-request-id: + - 97ca42c1-d8ce-4966-88e9-879d60847857 + method: POST + uri: https://llm.api.cloud.yandex.net/v1/chat/completions + response: + body: + string: 'data: {"id":"4ae6d707-cd09-4e3b-9622-9c79f5ae29a1","object":"chat.completion.chunk","created":1774968385,"model":"gpt://b1gr2p0etomug2vas4ie/yandexgpt/latest","choices":[{"index":0,"delta":{"role":"assistant","content":"Are"}}]} + + + data: {"id":"4ae6d707-cd09-4e3b-9622-9c79f5ae29a1","object":"chat.completion.chunk","created":1774968385,"model":"gpt://b1gr2p0etomug2vas4ie/yandexgpt/latest","choices":[{"index":0,"delta":{"content":" + you trying to tell me that you''re a cat? Or are you just in a playful mood? + Either way, meow back at you! Meow!"},"finish_reason":"stop"}]} + + + data: {"id":"4ae6d707-cd09-4e3b-9622-9c79f5ae29a1","object":"chat.completion.chunk","created":1774968385,"model":"gpt://b1gr2p0etomug2vas4ie/yandexgpt/latest","choices":[],"usage":{"prompt_tokens":29,"total_tokens":63,"completion_tokens":34,"prompt_tokens_details":{"cached_tokens":0}}} + + + data: [DONE] + + + ' + headers: + cache-control: + - no-cache + content-type: + - text/event-stream + date: + - Tue, 31 Mar 2026 14:46:25 GMT + server: + - ycalb + transfer-encoding: + - chunked + x-client-request-id: + - 97ca42c1-d8ce-4966-88e9-879d60847857 + x-request-id: + - 8f06983f-f3e7-480b-82b4-78dc0d5a2dfb + x-server-trace-id: + - 68a6526367f1c05b:94e2370f3b41800d:272414d44e47efff:1 + status: + code: 200 + message: OK +version: 1 diff --git a/tests/langchain_/test_chat_api.py b/tests/langchain_/test_chat_api.py new file mode 100644 index 00000000..6d2caf03 --- /dev/null +++ b/tests/langchain_/test_chat_api.py @@ -0,0 +1,40 @@ +"""Tests for LangChain integration via Chat API.""" + +from __future__ import annotations + +import pytest + +pytestmark = [pytest.mark.asyncio, pytest.mark.vcr, pytest.mark.require_env('langchain_core')] + + +@pytest.fixture(name='model') +def fixture_model(async_sdk): + return async_sdk.chat.completions('yandexgpt').langchain() + + +@pytest.fixture(name='chat_history') +def fixture_chat_history(): + from langchain_core.messages import AIMessage, HumanMessage # pylint: disable=import-outside-toplevel,import-error + + return [ + HumanMessage(content="hello!"), + AIMessage(content="Hi there human!"), + HumanMessage(content="Meow!"), + ] + + +async def test_ainvoke(model, chat_history): + result = await model.ainvoke(chat_history) + + assert result.content + assert result.usage_metadata is not None + assert result.usage_metadata["total_tokens"] > 0 + assert "finish_reason" in result.response_metadata + assert "model" in result.response_metadata + + +async def test_astream(model, chat_history): + chunks = [chunk async for chunk in model.astream(chat_history)] + + assert len(chunks) > 0 + assert any(c.content for c in chunks)