feat(types): add framework-agnostic LLM type system#1745
Conversation
Greptile SummaryThis PR introduces
|
| Filename | Overview |
|---|---|
| nemoguardrails/types.py | New framework-agnostic type system; clean implementation with correct validation, role aliases, and provider_metadata handling via extra-key capture. |
| tests/test_types.py | Comprehensive test suite covering all data types, Protocol conformance, roundtrip serialization, role aliases, and edge-case argument validation. |
Class Diagram
%%{init: {'theme': 'neutral'}}%%
classDiagram
class Role {
<<enumeration>>
USER = "user"
ASSISTANT = "assistant"
SYSTEM = "system"
TOOL = "tool"
}
class ToolCallFunction {
+str name
+Dict arguments
}
class ToolCall {
+str id
+str type
+ToolCallFunction function
+to_dict() Dict
}
class UsageInfo {
+int input_tokens
+int output_tokens
+int total_tokens
+Optional~int~ reasoning_tokens
+Optional~int~ cached_tokens
}
class ChatMessage {
+Role role
+Optional content
+Optional~List~ToolCall~~ tool_calls
+Optional~str~ tool_call_id
+Optional~str~ name
+Dict provider_metadata
+from_user(content) ChatMessage$
+from_assistant(content) ChatMessage$
+from_system(content) ChatMessage$
+from_tool(content, tool_call_id) ChatMessage$
+to_dict() Dict
+from_dict(d) ChatMessage$
}
class LLMResponse {
+str content
+Optional~str~ reasoning
+Optional~List~ToolCall~~ tool_calls
+Optional~str~ model
+Optional~FinishReason~ finish_reason
+Optional~UsageInfo~ usage
+Optional~Dict~ provider_metadata
}
class LLMResponseChunk {
+Optional~str~ delta_content
+Optional~str~ delta_reasoning
+Optional~List~ToolCall~~ delta_tool_calls
+Optional~FinishReason~ finish_reason
+Optional~UsageInfo~ usage
}
class LLMModel {
<<Protocol>>
+generate(prompt, stop, **kwargs) LLMResponse
+stream(prompt, stop, **kwargs) AsyncIterator
+model_name str
+provider_name Optional~str~
+provider_url Optional~str~
}
class LLMFramework {
<<Protocol>>
+create_model(model_name, provider_name, model_kwargs) LLMModel
}
ChatMessage --> Role
ChatMessage --> ToolCall
ToolCall --> ToolCallFunction
LLMResponse --> ToolCall
LLMResponse --> UsageInfo
LLMResponseChunk --> ToolCall
LLMResponseChunk --> UsageInfo
LLMModel --> LLMResponse : generates
LLMModel --> LLMResponseChunk : streams
LLMFramework --> LLMModel : creates
Reviews (9): Last reviewed commit: "final fixes" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
📝 WalkthroughWalkthroughIntroduced a new Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/test_types.py (1)
170-313: Add regression tests for unknown-key metadata capture and non-object JSON arguments.Please add assertions that:
- unknown top-level keys are merged into
provider_metadata, and- JSON arguments that decode to non-objects (e.g.,
"[]") raiseValueError.
These lock the criticalfrom_dictedge behavior.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_types.py` around lines 170 - 313, Add two regression assertions to the test suite for ChatMessage.from_dict: (1) call ChatMessage.from_dict with an extra unknown top-level key (e.g., "unexpected_key": "v") and assert the resulting msg.provider_metadata contains that key/value (reference ChatMessage.from_dict and provider_metadata); (2) add a test case where a tool_call.function.arguments is a JSON string that decodes to a non-object (e.g., "[]") and assert ChatMessage.from_dict raises ValueError (reference ChatMessage.from_dict and ToolCall.function.arguments). Ensure these are added near the existing tool_calls/arguments tests so they run alongside related cases.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nemoguardrails/types.py`:
- Around line 149-155: After deserializing raw_args (via json.loads) into
args_dict, ensure the result is a mapping/dict and not a JSON array/primitive;
if args_dict is not an instance of dict (e.g., list, int, str), raise a
ValueError indicating "Tool call arguments must be a JSON object" and include
the offending raw_args; update the branch handling raw_args (referencing
raw_args, args_dict, json.loads) to perform this type check and error raise so
downstream adapters only receive Dict[str, Any].
- Around line 121-175: from_dict in ChatMessage currently ignores unknown
top-level keys instead of preserving them in provider_metadata; update
ChatMessage.from_dict to collect all keys from the input dict `d` that are not
the known keys ("role", "content", "tool_calls", "tool_call_id", "name",
"provider_metadata") and merge them into the returned provider_metadata (also
merging any existing d.get("provider_metadata") dict), so provider-specific
fields are preserved. Keep existing behavior for parsing role, tool_calls
(ToolCall/ToolCallFunction), and JSON argument decoding, but ensure the final
cls(...) call uses the merged provider_metadata dict.
---
Nitpick comments:
In `@tests/test_types.py`:
- Around line 170-313: Add two regression assertions to the test suite for
ChatMessage.from_dict: (1) call ChatMessage.from_dict with an extra unknown
top-level key (e.g., "unexpected_key": "v") and assert the resulting
msg.provider_metadata contains that key/value (reference ChatMessage.from_dict
and provider_metadata); (2) add a test case where a tool_call.function.arguments
is a JSON string that decodes to a non-object (e.g., "[]") and assert
ChatMessage.from_dict raises ValueError (reference ChatMessage.from_dict and
ToolCall.function.arguments). Ensure these are added near the existing
tool_calls/arguments tests so they run alongside related cases.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 33cc0dc4-9b96-4c75-9b0f-8edfa02c00c0
📒 Files selected for processing (2)
nemoguardrails/types.pytests/test_types.py
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Pouyanpi
force-pushed
the
feat/langchain-decouple/stack-1-canonical-types
branch
from
4bb2ddf to
8f7e8d9
Compare
Introduce nemoguardrails/types.py with provider-independent data types and protocols that replace direct LangChain type dependencies in core code.
Pouyanpi
force-pushed
the
feat/langchain-decouple/stack-1-canonical-types
branch
from
8563b47 to
320d754
Compare
Only LangChain needs mode (chat vs text). Other frameworks ignore it. Move mode into model_kwargs so each framework extracts what it needs.
|
thanks @tgasser-nv, all the comments are addressed/resolved. For the tool.id i've opened this PR: #1775 |
tgasser-nv
left a comment
tgasser-nv
left a comment
There was a problem hiding this comment.
Looks good. Only concern is the impedance mismatch between the Model pydantic object and provider_name, prodivder_url. Why do we need a second slightly different way of representing the same thing? Are you intending to deprecate Model in the future?
|
Thanks @tgasser-nv . re your question I see it this way: provider_name/model_name/provider_url are explicit for someone implementing an adapter and conform to our existing internals: llm_call() parameters (model_name, model_provider), LLMCallInfo fields (llm_model_name, llm_provider_name), LLMCallException error context (model=, provider=, endpoint=), and OpenTelemetry tracing spans (GenAIAttributes.GEN_AI_PROVIDER_NAME). The protocol surfaces what the codebase already names this way. The mapping happens in exactly one place (llmrails.py, When you review the other PRs in the stack it becomes clearer how these fit together, but if it's still a point of debate we can rename in a follow-up PR before the release. |
tgasser-nv
left a comment
tgasser-nv
left a comment
There was a problem hiding this comment.
Looks good! Please take a look at the feedback, only a few small cleanups needed before meging
|
a final note about engine vs provider: we can have openai |
2 participants
Part of the LangChain decoupling stack:
Description
Introduce nemoguardrails/types.py with provider-independent data types and protocols that replace direct LangChain type dependencies in core code.
nemoguardrails/types.pywith framework-agnostic data types (ChatMessage,Role,ToolCall,ToolCallFunction,LLMResponse,LLMResponseChunk,UsageInfo,FinishReason) and protocols (LLMModel,LLMFramework)ChatMessage.from_dict()handles both OpenAI nested and legacy flat tool call formats, JSON string argument parsing, and role aliases (bot,human,developer)ChatMessage.to_dict()/from_dict()roundtrip preserves all fields includingprovider_metadataLLMModelprotocol defines thegenerate()/stream()contract that all LLM adapters must implementLLMFrameworkprotocol defines thecreate_model()factory contract for pluggable backends