chore(weave): custom provider model built in object backend #3841
Conversation
|
Preview this PR with FeatureBee: https://beta.wandb.ai/?betaVersion=689eac05076100269778c89a4a497f494b8792b6 |
WalkthroughThis pull request introduces several new Zod schemas and TypeScript types for built-in object classes related to language model configurations and providers. It updates error handling and control flow in the trace server by distinguishing between standard and custom models. New classes and enumerations are added in separate files for LLM models and providers, and these classes are registered in the global registry. Additionally, the custom provider flow is enhanced by updating the API call parameters and including a dedicated function to extract provider details. Changes
Sequence Diagram(s)sequenceDiagram
participant C as Client Request
participant CS as ClickHouseTraceServer
participant CP as CustomProviderService
participant LLM as lite_llm_completion
C->>CS: Send completions_create request
CS->>CS: Check model info (standard vs custom)
alt Custom Model
CS->>CP: get_custom_provider_info(project_id, model_name, obj_read_func)
CP-->>CS: Return provider details or error
CS->>LLM: Call lite_llm_completion(api_key, inputs, provider, base_url, extra_headers, return_type)
else Standard Model
CS->>LLM: Call lite_llm_completion(api_key, inputs, provider)
end
LLM-->>CS: Return API response
CS-->>C: Send final response
Suggested reviewers
Poem
✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (9)
weave/trace_server/interface/builtin_object_classes/provider.py (2)
6-7: Consider adding more ReturnType options for future extensibility.Currently, the
ReturnTypeenum only includesOPENAI. As you're building custom provider functionality, you might want to consider whether other return types (likeANTHROPIC,COHERE, etc.) should be added to make the enum more future-proof.
1-14: Add docstrings to explain class purposes and usage.The file lacks explanatory docstrings that would help developers understand the purpose of these classes and how to use them. Consider adding class-level docstrings explaining what a
Providerrepresents in your system and how theReturnTypeimpacts functionality.from enum import Enum from weave.trace_server.interface.builtin_object_classes import base_object_def class ReturnType(str, Enum): + """Defines the format in which provider responses should be returned. + + Currently only supports OpenAI-compatible response format. + """ OPENAI = "openai" class Provider(base_object_def.BaseObject): + """Represents an LLM API provider configuration. + + Stores connection details for an LLM provider API including the base URL, + the name of the API key to use, any extra headers required, and the expected + response format. + """ base_url: str api_key_name: str extra_headers: dict[str, str] = {} return_type: ReturnType = ReturnType.OPENAIweave/trace_server/interface/builtin_object_classes/llm_model.py (1)
1-33: Add docstrings to explain relationships between model classes.The file lacks explanatory docstrings that would help developers understand the purpose of these classes and their relationships. Clear documentation would help explain the difference between
LLMModelandProviderModeland how they work together.from enum import Enum from typing import Optional, Union from pydantic import BaseModel from weave.trace_server.interface.builtin_object_classes import base_object_def class ModelMode(str, Enum): + """Defines the interaction mode for language models. + + COMPLETION: For models that complete text (continuation) + CHAT: For models that engage in conversational exchanges + """ COMPLETION = "completion" CHAT = "chat" class ModelParams(BaseModel): + """Configuration parameters for language model inference. + + Contains common parameters used to control language model outputs + such as temperature, top_p sampling, token limits, and penalties. + """ temperature: Optional[float] = None top_p: Optional[float] = None max_tokens: Optional[int] = None presence_penalty: Optional[float] = None frequency_penalty: Optional[float] = None stop: Optional[Union[str, list[str]]] = None class ProviderModel(base_object_def.BaseObject): + """Represents a specific model offered by a provider. + + Links to a provider configuration and specifies model-specific + details such as maximum token length and interaction mode. + """ provider: base_object_def.RefStr max_tokens: int mode: ModelMode = ModelMode.CHAT class LLMModel(base_object_def.BaseObject): + """Represents a configured LLM for use in the application. + + Connects to a provider model and optionally includes a default prompt + and parameter settings to use with the model. + """ provider_model: base_object_def.RefStr prompt: Optional[base_object_def.RefStr] = None default_params: ModelParams = ModelParams()weave/trace_server/clickhouse_trace_server_batched.py (2)
1666-1679: Improve error handling granularity.The current error handling uses a broad try-except block that catches all exceptions, which could mask important specific errors. Consider catching more specific exceptions or at least logging the exception details before returning the error response.
try: ( base_url, api_key, extra_headers, return_type, actual_model_name, ) = get_custom_provider_info( project_id=req.project_id, model_name=model_name, obj_read_func=self.obj_read, ) -except Exception as e: +except NotFoundError as e: return tsi.CompletionsCreateRes(response={"error": str(e)}) +except ValueError as e: + # Handle specific validation errors + logger.warning(f"Invalid custom provider configuration: {e}") + return tsi.CompletionsCreateRes(response={"error": f"Invalid provider configuration: {str(e)}"}) +except Exception as e: + # Log unexpected errors + logger.error(f"Unexpected error in custom provider flow: {e}", exc_info=True) + return tsi.CompletionsCreateRes(response={"error": f"Internal error: {str(e)}"})
1630-1661: Add validation for API key existence before proceeding.For better error handling, consider validating that the API key exists before proceeding with the API call, especially for providers that aren't handled like bedrock.
# Handle standard model case # 1. We get the model info from the map # 2. We fetch the API key, with the secret fetcher # 3. We set the provider, to the litellm provider # 4. If no api key, we raise an error, except for bedrock and bedrock_converse (we fetch bedrock credentials, in lite_llm_completion) secret_name = model_info.get("api_key_name") if not secret_name: raise InvalidRequest(f"No secret name found for model {model_name}") secret_fetcher = _secret_fetcher_context.get() if not secret_fetcher: raise InvalidRequest( f"No secret fetcher found, cannot fetch API key for model {model_name}" ) api_key = ( secret_fetcher.fetch(secret_name).get("secrets", {}).get(secret_name) ) provider = model_info.get("litellm_provider", "openai") # We fetch bedrock credentials, in lite_llm_completion, later if not api_key and provider != "bedrock" and provider != "bedrock_converse": raise MissingLLMApiKeyError( f"No API key {secret_name} found for model {model_name}", api_key_name=secret_name, ) + +# Log provider information for debugging +logger.debug(f"Using provider {provider} for model {model_name}")weave/trace_server/llm_completion.py (2)
42-66: Catch narrower exceptions or provide additional error handling insights
Catching a broadExceptionhere can obscure specific failure modes. Consider narrowing the scope or at least logging more debug information (like full tracebacks) before returning the error response. Additionally, evaluating or sanitizing the user-suppliedbase_urlcan help avoid potential security vulnerabilities.
155-259: Validate the exact composition of the model name parts
This function only checks if themodel_namesplit length is< 2; consider verifying precisely two parts (e.g., “provider_id/model_id”) to avoid silently ignoring extra segments. This small adjustment can help maintain data integrity and reduce user confusion.weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/generatedBuiltinObjectClasses.zod.ts (2)
9-10: Future-proofing return types
Currently,'openai'is the sole enum value forReturnTypeSchema. If future expansions are planned (e.g.,'custom','bedrock'), consider adding them here to avoid migrations later.
41-49: Unify the type of “stop” parameter
Allowing'stop'to be a single string or array may complicate downstream usage. Converting single strings into a one-element array can simplify internal logic and reduce condition checks.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
weave/trace_server/interface/builtin_object_classes/generated/generated_builtin_object_class_schemas.jsonis excluded by!**/generated/**
📒 Files selected for processing (6)
weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/generatedBuiltinObjectClasses.zod.ts(4 hunks)weave/trace_server/clickhouse_trace_server_batched.py(3 hunks)weave/trace_server/interface/builtin_object_classes/builtin_object_registry.py(2 hunks)weave/trace_server/interface/builtin_object_classes/llm_model.py(1 hunks)weave/trace_server/interface/builtin_object_classes/provider.py(1 hunks)weave/trace_server/llm_completion.py(4 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
`**/*.py`: Focus on pythonic code patterns. Check for proper...
**/*.py: Focus on pythonic code patterns.
Check for proper exception handling.
Verify type hints usage where applicable.
Look for potential performance improvements.
Don't comment on formatting if black/isort is configured.
Check for proper dependency injection patterns.
Verify proper async handling if applicable.
weave/trace_server/interface/builtin_object_classes/provider.pyweave/trace_server/interface/builtin_object_classes/builtin_object_registry.pyweave/trace_server/clickhouse_trace_server_batched.pyweave/trace_server/llm_completion.pyweave/trace_server/interface/builtin_object_classes/llm_model.py
`**/*.{js,jsx,ts,tsx}`: Focus on architectural and logical i...
**/*.{js,jsx,ts,tsx}: Focus on architectural and logical issues rather than style (assuming ESLint is in place).
Flag potential memory leaks and performance bottlenecks.
Check for proper error handling and async/await usage.
Avoid strict enforcement of try/catch blocks - accept Promise chains, early returns, and other clear error handling patterns. These are acceptable as long as they maintain clarity and predictability.
Ensure proper type usage in TypeScript files.
Look for security vulnerabilities in data handling.
Don't comment on formatting if prettier is configured.
Verify proper React hooks usage and component lifecycle.
Check for proper state management patterns.
weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/generatedBuiltinObjectClasses.zod.ts
🔇 Additional comments (16)
weave/trace_server/interface/builtin_object_classes/builtin_object_registry.py (2)
9-12: LGTM! Imports are correctly organized.The imports for the new classes are properly organized and follow the existing import pattern.
37-39: LGTM! Classes are properly registered.The new classes are correctly registered in the registry using the established pattern.
weave/trace_server/interface/builtin_object_classes/llm_model.py (4)
9-11: LGTM! ModelMode enum is properly defined.The enum is correctly implemented as a string-based enum with appropriate values for completion and chat modes.
14-20: LGTM! ModelParams are properly defined with appropriate types.The parameters cover common LLM configuration options with correct optional typing.
23-26: LGTM! ProviderModel is properly implemented.The class correctly extends BaseObject and includes the necessary fields with appropriate types.
29-32: LGTM! LLMModel is properly implemented.The class correctly extends BaseObject and includes the necessary fields with appropriate types.
weave/trace_server/clickhouse_trace_server_batched.py (4)
89-92: LGTM! Import statements properly organized.The imports for
get_custom_provider_infoandlite_llm_completionare correctly added.
1616-1685: Clear organization of standard vs. custom model handling.The code clearly organizes the flow for handling standard models versus custom models with good inline comments explaining the logic.
1690-1697: LGTM! API call parameters properly expanded.The
lite_llm_completioncall correctly uses keyword arguments with all the necessary parameters for both standard and custom model flows.
1721-1721: LGTM! Model name variable usage is consistent.The code correctly uses the
model_namevariable instead of directly referencingreq.inputs.modelfor consistency.weave/trace_server/llm_completion.py (1)
17-19: New parameters appear consistent and well-typed
The optional parameters forbase_url,extra_headers, andreturn_typeintegrate smoothly with the rest of the function and follow Pythonic type hint conventions.weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/generatedBuiltinObjectClasses.zod.ts (5)
12-13: Model mode definitions look good
TheModelModeSchemawith'chat'and'completion'is straightforward and aligns with typical usage.
51-60: Provider schema aligns with custom provider approach
No issues found. Requiring abase_urlandapi_key_nameis consistent with the addition of custom provider logic.
61-69: Clear definition of provider models
Mandatorymax_tokensand optionalmodefields provide flexibility. The schema clearly enforces minimal constraints for validProviderModelobjects.
100-108: LLM model structure is well-formed
This expanded schema correctly referencesdefault_paramsand retains optional fields. It looks coherent with new functionality around providers.
122-124: Successful registry updates
AddingLLMModel,Provider, andProviderModelto the registry ensures these objects are recognized throughout the application.
Description
Adds 3 built in objects
provider
provider_model
LLM_model
adds handling in clickhouse_trace_server_batched and llm_completion for custom models
Diagram of how the new objects interact
video of it working (for this i spun up a server on my local that just passes stuff through to openai) caveat for this to work with litellm, the apis need to be the same as openai
Screen.Recording.2025-03-04.at.12.29.16.PM.mov
core branch with frontend - core/custom-beta
Testing
How was this PR tested?
Summary by CodeRabbit
New Features
Refactor