Skip to content

Use Async streaming by default #25

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anuj-tambwekar wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Use Async streaming by default #25

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 83 additions & 39 deletions langchain_gradient/chat_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,23 @@
import json
import logging
import os
from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Union
import asyncio
from typing import (
Any,
Callable,
Dict,
Iterator,
AsyncIterator,
List,
Optional,
Sequence,
Union,
)

from gradient import Gradient
from gradient import Gradient, AsyncGradient
from langchain_core.callbacks import (
CallbackManagerForLLMRun,
AsyncCallbackManagerForLLMRun,
)
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import (
Expand Down Expand Up @@ -108,6 +120,7 @@ class ChatGradient(BaseChatModel):
_stream(messages, ...)
Stream chat completions for the given messages.
"""

api_key: Optional[str] = Field(
default=os.environ.get("DIGITALOCEAN_INFERENCE_KEY"),
exclude=True,
Expand Down Expand Up @@ -167,12 +180,12 @@ def validate_temperature(cls, values: dict[str, Any]) -> Any:

@property
def user_agent_package(self) -> str:
return f"LangChain"
return "LangChain"

@property
def user_agent_version(self) -> str:
return "0.1.22"

@property
def _llm_type(self) -> str:
"""Return type of chat model."""
Expand Down Expand Up @@ -210,17 +223,20 @@ def _update_parameters_with_model_fields(self, parameters: dict) -> None:
def _convert_message(self, msg: BaseMessage) -> Dict[str, Any]:
"""Convert a LangChain message to OpenAI API format."""
if hasattr(msg, "type"):
role = {"human": "user", "ai": "assistant", "system": "system", "tool": "tool"}.get(
msg.type, msg.type
)
role = {
"human": "user",
"ai": "assistant",
"system": "system",
"tool": "tool",
}.get(msg.type, msg.type)
else:
role = getattr(msg, "role", "user")

result: Dict[str, Any] = {"role": role, "content": msg.content}

if isinstance(msg, ToolMessage):
result["tool_call_id"] = msg.tool_call_id

if isinstance(msg, AIMessage) and msg.tool_calls:
result["tool_calls"] = [
{
Expand All @@ -233,7 +249,7 @@ def _convert_message(self, msg: BaseMessage) -> Dict[str, Any]:
}
for tc in msg.tool_calls
]

return result

def _parse_tool_calls(self, raw_tool_calls: List[Any]) -> List[ToolCall]:
Expand All @@ -245,15 +261,21 @@ def _parse_tool_calls(self, raw_tool_calls: List[Any]) -> List[ToolCall]:
if hasattr(tc, "function"):
func = tc.function
name = func.name if hasattr(func, "name") else func.get("name")
arguments = func.arguments if hasattr(func, "arguments") else func.get("arguments")
arguments = (
func.arguments
if hasattr(func, "arguments")
else func.get("arguments")
)
tc_id = tc.id if hasattr(tc, "id") else tc.get("id")
else:
func = tc.get("function", {})
name = func.get("name")
arguments = func.get("arguments")
tc_id = tc.get("id")

args = json.loads(arguments) if isinstance(arguments, str) else arguments

args = (
json.loads(arguments) if isinstance(arguments, str) else arguments
)
tool_calls.append(create_tool_call(name=name, args=args, id=tc_id))
except (json.JSONDecodeError, KeyError, AttributeError) as e:
logger.warning(
Expand Down Expand Up @@ -293,7 +315,7 @@ def _generate(

tools = kwargs.get("tools") or self.tools
tool_choice = kwargs.get("tool_choice") or self.tool_choice

if tools:
parameters["tools"] = tools
if tool_choice:
Expand All @@ -306,7 +328,7 @@ def _generate(
if hasattr(choice.message, "content")
else choice.message
) or "" # Ensure content is never None

usage = getattr(completion, "usage", {})
response_metadata = {
"finish_reason": getattr(choice, "finish_reason", None),
Expand All @@ -318,21 +340,21 @@ def _generate(
"model_name": getattr(completion, "model", None),
"id": getattr(completion, "id", None),
}

tool_calls = []
raw_tool_calls = getattr(choice.message, "tool_calls", None)
if raw_tool_calls:
tool_calls = self._parse_tool_calls(raw_tool_calls)

message_kwargs: Dict[str, Any] = {
"content": content,
"additional_kwargs": {"refusal": getattr(choice.message, "refusal", None)},
"response_metadata": response_metadata,
}

if tool_calls:
message_kwargs["tool_calls"] = tool_calls

if self.stream_options and self.stream_options.get("include_usage"):
message_kwargs["usage_metadata"] = {
"input_tokens": getattr(usage, "prompt_tokens", None),
Expand All @@ -343,49 +365,50 @@ def _generate(
generation = ChatGeneration(message=message)
return ChatResult(generations=[generation])

def _stream(
async def _astream(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> Iterator[ChatGenerationChunk]:
) -> AsyncIterator[ChatGenerationChunk]:
if not self.api_key:
raise ValueError(
"Gradient model access key not provided. Set DIGITALOCEAN_INFERENCE_KEY env var or pass api_key param."
)

inference_client = Gradient(
# Use the async client for streaming to prevent buffering issues
inference_client = AsyncGradient(
model_access_key=self.api_key,
base_url="https://inference.do-ai.run/v1",
user_agent_package=self.user_agent_package,
user_agent_version=self.user_agent_version,
user_agent_version=self.user_agent_version,
)

parameters: Dict[str, Any] = {
"messages": [self._convert_message(m) for m in messages],
"stream": True, # Enable streaming
"model": self.model_name,
}

self._update_parameters_with_model_fields(parameters)

# Handle tools from kwargs (e.g., from bind_tools)
tools = kwargs.get("tools") or self.tools
tool_choice = kwargs.get("tool_choice") or self.tool_choice

if tools:
parameters["tools"] = tools
if tool_choice:
parameters["tool_choice"] = tool_choice

try:
stream = inference_client.chat.completions.create(**parameters)
for completion in stream:
stream = await inference_client.chat.completions.create(**parameters)
async for completion in stream:
delta = completion.choices[0].delta

content = getattr(delta, "content", None) or ""

tool_call_chunks = []
raw_tool_calls = getattr(delta, "tool_calls", None)
if raw_tool_calls:
Expand All @@ -402,7 +425,7 @@ def _stream(
arguments = func.get("arguments")
tc_id = tc.get("id")
index = tc.get("index")

tool_call_chunks.append(
create_tool_call_chunk(
name=name,
Expand All @@ -411,17 +434,17 @@ def _stream(
index=index,
)
)

if not content and not tool_call_chunks:
continue

chunk_kwargs: Dict[str, Any] = {"content": content}
if tool_call_chunks:
chunk_kwargs["tool_call_chunks"] = tool_call_chunks

chunk = ChatGenerationChunk(message=AIMessageChunk(**chunk_kwargs))
if run_manager and content:
run_manager.on_llm_new_token(content, chunk=chunk)
await run_manager.on_llm_new_token(content, chunk=chunk)
yield chunk

# Optionally yield usage metadata at the end if available
Expand All @@ -448,6 +471,27 @@ def _stream(
)
yield error_chunk

def _stream(
self,
*args: Any,
**kwargs: Any,
) -> Iterator[ChatGenerationChunk]:
loop = asyncio.get_event_loop()
queue: asyncio.Queue = asyncio.Queue()

async def producer() -> None:
async for chunk in self._astream(*args, **kwargs):
await queue.put(chunk)
await queue.put(None)

loop.create_task(producer())

while True:
item = loop.run_until_complete(queue.get())
if item is None:
break
yield item

@property
def init_from_env_params(self) -> tuple[dict, dict, dict]:
# env_vars, model_params, expected_attrs
Expand Down Expand Up @@ -501,19 +545,19 @@ def bind_tools(
... def get_weather(location: str) -> str:
... '''Get weather for a location.'''
... return f"Weather in {location}: sunny"
>>>
>>>
>>> llm = ChatGradient(model="llama3.3-70b-instruct")
>>> llm_with_tools = llm.bind_tools([get_weather])
>>> response = llm_with_tools.invoke("What's the weather in SF?")
"""
formatted_tools = [convert_to_openai_tool(tool) for tool in tools]

formatted_tool_choice = tool_choice
if tool_choice == "any":
formatted_tool_choice = "required"
elif tool_choice is None:
formatted_tool_choice = "auto"

return self.bind(
tools=formatted_tools,
tool_choice=formatted_tool_choice,
Expand Down