feat: add support langchain astream
#936
Merged
+106
−1
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Contributor
Reviewer's GuideThis PR implements support for langchain's astream streaming by introducing a token accumulation mechanism and flushing complete responses into spans, adds a new chain-level error callback with enhanced error classification and cleanup logic, and bumps the SDK version to 1.36.1. Sequence diagram for streaming token accumulation and response flushing in LangChain astreamsequenceDiagram
participant LangChain
participant CallbackHandler
participant Span
LangChain->>CallbackHandler: on_llm_new_token(token, run_id)
CallbackHandler->>CallbackHandler: Accumulate token in streaming_chunks[run_id]
LangChain->>CallbackHandler: on_llm_end(response, run_id)
CallbackHandler->>CallbackHandler: Join tokens in streaming_chunks[run_id]
CallbackHandler->>Span: set_attribute(GEN_AI_CONTENT_COMPLETION, complete_response)
CallbackHandler->>Span: set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
CallbackHandler->>CallbackHandler: Clean up streaming_chunks[run_id]
CallbackHandler->>Span: _process_llm_response(span, response, run_id)
Sequence diagram for chain-level error handling and cleanupsequenceDiagram
participant LangChain
participant CallbackHandler
participant Span
LangChain->>CallbackHandler: on_chain_error(error, run_id)
CallbackHandler->>CallbackHandler: Clean up streaming_chunks[run_id]
CallbackHandler->>Span: set_attribute(GEN_AI_FRAMEWORK_ERROR_CLASS, error_class)
CallbackHandler->>Span: set_attribute(GEN_AI_FRAMEWORK_ERROR_TYPE, type(error).__name__)
CallbackHandler->>Span: set_attribute(GEN_AI_FRAMEWORK_ERROR_MESSAGE, str(error))
CallbackHandler->>Span: set_status(StatusCode.ERROR, str(error))
CallbackHandler->>Span: record_exception(error)
CallbackHandler->>CallbackHandler: _end_span(run_id)
Class diagram for updated CallbackHandler streaming and error handlingclassDiagram
class CallbackHandler {
spans: Dict[UUID, SpanHolder]
streaming_chunks: Dict[UUID, List[str]]
session_name: str
tags_enabled: bool
events_enabled: bool
on_llm_new_token(token: str, chunk: Optional[Any], run_id: UUID, parent_run_id: Optional[UUID], **kwargs)
on_llm_end(response: Any, run_id: UUID, parent_run_id: Optional[UUID], **kwargs)
on_llm_error(error: Exception, run_id: UUID, parent_run_id: Optional[UUID], **kwargs)
on_chain_error(error: Exception, run_id: UUID, parent_run_id: Optional[UUID], **kwargs)
}
CallbackHandler "1" o-- "*" SpanHolder
CallbackHandler "1" o-- "*" streaming_chunks
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Contributor
There was a problem hiding this comment.
Hey there - I've reviewed your changes - here's some feedback:
- You’re duplicating streaming_chunks cleanup in on_llm_end, on_llm_error, and on_chain_error—extract that into a shared helper to reduce repetition.
- If callbacks can run concurrently, unsynchronized access to streaming_chunks may race—consider guard (e.g. a lock) or a thread-safe structure.
- Accumulating all tokens unbounded could exhaust memory for long streams—think about capping buffer size or flushing partial results periodically.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- You’re duplicating streaming_chunks cleanup in on_llm_end, on_llm_error, and on_chain_error—extract that into a shared helper to reduce repetition.
- If callbacks can run concurrently, unsynchronized access to streaming_chunks may race—consider guard (e.g. a lock) or a thread-safe structure.
- Accumulating all tokens unbounded could exhaust memory for long streams—think about capping buffer size or flushing partial results periodically.
## Individual Comments
### Comment 1
<location> `sdk/python/src/openlit/instrumentation/langchain/callback_handler.py:663-665` </location>
<code_context>
+
+ self._end_span(run_id)
+
+ except Exception:
+ # Graceful error handling
+ pass
+
def on_chat_model_start(
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Swallowing all exceptions may obscure underlying issues.
Instead of passing silently, log or record the exception to aid diagnostics unless you have a specific reason to suppress all errors.
```suggestion
except Exception as exc:
# Log the exception for diagnostics
import logging
logging.exception("Exception occurred while ending span for run_id %s", run_id)
```
</issue_to_address>
### Comment 2
<location> `sdk/python/src/openlit/instrumentation/langchain/callback_handler.py:801` </location>
<code_context>
+ self.streaming_chunks[run_id] = []
+
+ # Accumulate the token
+ if token:
+ self.streaming_chunks[run_id].append(token)
+
</code_context>
<issue_to_address>
**question:** Empty string tokens are ignored; clarify if this is intentional.
If empty tokens are valid from the source, skipping them could cause data loss. Please confirm this is the intended behavior.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Important
feat: ...orfix: ....Issue number: resolves #933
Change description:
Checklist
If your change doesn't seem to apply, please leave them unchecked.
feat: ...orfix: ....Acknowledgment
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license.
Summary by Sourcery
Add support for LangChain
astreamstreaming by tracking individual tokens via on_llm_new_token, assembling full streamed responses in on_llm_end, and capturing usage metrics. Enhance error handling by cleaning up streaming buffers and classifying both LLM and chain errors with detailed attributes. Bump Python SDK version to 1.36.1.New Features:
Enhancements:
Build: