Skip to content

Change session_metadata_api Session from std lib dataclass to Pydantic #313

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

Merged
merged 4 commits into from
Oct 7, 2025
Merged

Change session_metadata_api Session from std lib dataclass to Pydantic #313

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8f3f364
Convert Session into Pydantic
mliu-cloudera Sep 2, 2025
db8fb68
Remove creation date and username
mliu-cloudera Sep 2, 2025
9f1a589
Enable field validation
mliu-cloudera Sep 2, 2025
8f19395
Swap lines to make a bit more sequential sense
mliu-cloudera Sep 2, 2025
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
46 changes: 23 additions & 23 deletions llm-service/app/services/metadata_apis/session_metadata_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,43 +35,49 @@
# BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
# DATA.
#
import dataclasses
Copy link

Copilot AI Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] The dataclasses import is only used for the UpdatableSession class. Consider removing this import and converting UpdatableSession to Pydantic as well for consistency, or add a comment explaining why UpdatableSession remains a dataclass.

Copilot uses AI. Check for mistakes.
Copy link
Collaborator Author

@mliu-cloudera mliu-cloudera Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the time it would take (to be fair, minutes) to test changing UpdatableSession too isn't too worth it right now. But maybe I'm being too lazy 😶

import json
from dataclasses import dataclass, field
from datetime import datetime
from typing import List, Any, Optional
Comment on lines +38 to 40
Copy link

Copilot AI Oct 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The dataclasses import is only used for the UpdatableSession class. Consider moving this import closer to where it's used or removing it if UpdatableSession will also be migrated to Pydantic.

Copilot uses AI. Check for mistakes.

import requests
from pydantic import BaseModel, ConfigDict, Field, alias_generators

from app.config import settings
from app.services.utils import raise_for_http_error, body_to_json


@dataclass
class SessionQueryConfiguration:
class SessionQueryConfiguration(BaseModel):
model_config = ConfigDict(
alias_generator=alias_generators.to_camel,
validate_by_name=True,
revalidate_instances="always",
)

enable_hyde: bool
enable_summary_filter: bool
enable_tool_calling: bool = False
selected_tools: list[str] = field(default_factory=list)
selected_tools: list[str] = Field(default_factory=list)
disable_streaming: bool = False


@dataclass
class Session:
class Session(BaseModel):
model_config = ConfigDict(
alias_generator=alias_generators.to_camel,
Copy link
Collaborator Author

@mliu-cloudera mliu-cloudera Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This allows us to output a camelCased dictionary by calling .model_dump(by_alias=True)!

validate_by_name=True,
Copy link
Collaborator Author

@mliu-cloudera mliu-cloudera Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is needed to allow fields to still be set by their snake_case names.

revalidate_instances="always",
Copy link
Collaborator Author

@mliu-cloudera mliu-cloudera Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Without this, Session.model_validate(session) is little more than a fancy instanceof() check. With it, the types of fields are actually validated.

Probably not necessary everywhere we have Pydantic models at this time, but I need it here for a check in #307.

)

id: int
name: str
data_source_ids: List[int]
data_source_ids: list[int]
project_id: int
time_created: datetime
time_updated: datetime
created_by_id: str
updated_by_id: str
Comment on lines -64 to -67
Copy link
Collaborator Author

@mliu-cloudera mliu-cloudera Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing these because time_created and time_updated were causing a slight pain serializing to and from JSON (also turns out they were floats this whole time), and it doesn't seem like we're actually using these fields anywhere.

Copy link
Collaborator

@baasitsharief baasitsharief Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does removing these raise an error while updating sessions? or it just works fine?

Copy link
Collaborator Author

@mliu-cloudera mliu-cloudera Oct 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great question! I've tested this locally, and encounter no errors when updating a session, nor creating a new one (which calls /rename-session), nor chatting.

We're not actually using these fields when making session update requests anyway (since they're managed at the DB layer), and we're also not using them here in the Python service.

inference_model: str
rerank_model: str
rerank_model: Optional[str]
Copy link
Collaborator Author

@mliu-cloudera mliu-cloudera Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems rerank_model can come back as None from the Java backend; it was never an error before because dataclasses.dataclass doesn't validate types.

response_chunks: int
query_configuration: SessionQueryConfiguration
associated_data_source_id: Optional[int] = None

def get_all_data_source_ids(self) -> List[int]:
def get_all_data_source_ids(self) -> list[int]:
"""
Returns all data source IDs associated with the session.
If the session has an associated data source ID, it is included in the list.
Expand All @@ -81,14 +87,14 @@ def get_all_data_source_ids(self) -> List[int]:
)


@dataclass
@dataclasses.dataclass
class UpdatableSession:
id: int
name: str
dataSourceIds: List[int]
projectId: int
inferenceModel: str
rerankModel: str
rerankModel: Optional[str]
responseChunks: int
queryConfiguration: dict[str, bool | List[str]]
associatedDataSourceId: Optional[int]
Expand All @@ -114,10 +120,6 @@ def session_from_java_response(data: dict[str, Any]) -> Session:
name=data["name"],
data_source_ids=data["dataSourceIds"],
project_id=data["projectId"],
time_created=datetime.fromtimestamp(data["timeCreated"]),
time_updated=datetime.fromtimestamp(data["timeUpdated"]),
created_by_id=data["createdById"],
updated_by_id=data["updatedById"],
inference_model=data["inferenceModel"],
rerank_model=data["rerankModel"],
response_chunks=data["responseChunks"],
Expand All @@ -127,9 +129,7 @@ def session_from_java_response(data: dict[str, Any]) -> Session:
enable_tool_calling=data["queryConfiguration"].get(
"enableToolCalling", False
),
disable_streaming=data["queryConfiguration"].get(
"disableStreaming", False
),
disable_streaming=data["queryConfiguration"].get("disableStreaming", False),
Copy link

Copilot AI Oct 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] This line formatting change appears to be unrelated to the Pydantic migration. Consider keeping the original multi-line format for consistency with the surrounding code style.

Copilot uses AI. Check for mistakes.
selected_tools=data["queryConfiguration"]["selectedTools"] or [],
),
associated_data_source_id=data.get("associatedDataSourceId", None),
Expand Down
Loading