LCORE-1051: final updates

[tisnik]

Description

LCORE-1051: final updates

Type of change

• Refactor
• New feature
• Bug fix
• CVE fix
• Optimization
• Documentation Update
• Configuration Update
• Bump-up service version
• Bump-up dependent library
• Bump-up library or tool used for development (does not change the final image)
• CI configuration change
• Konflux configuration change
• Unit tests improvement
• Integration tests improvement
• End to end tests improvement

Tools used to create PR

• Assisted-by: N/A
• Generated by: N/A

Related Tickets & Documents

• Related Issue #LCORE-1051

Summary by CodeRabbit

• Refactor
  • Internal improvements to code quality and type safety across API endpoints and utility modules for enhanced maintainability and developer experience.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Walkthrough

This pull request standardizes type annotations across endpoint and utility modules by replacing union-with-None syntax (e.g., str | None) with the Optional[...] form from the typing module. Changes affect internal state declarations, function signatures, property return types, and method parameters across seven files without modifying runtime logic or behavior.

Changes

Cohort / File(s)	Summary
Endpoint module type annotation standardization src/app/endpoints/a2a.py, src/app/endpoints/query.py, src/app/endpoints/query_v2.py, src/app/endpoints/streaming_query.py, src/app/endpoints/streaming_query_v2.py	Converted union-with-None type hints to Optional[...] across public method signatures and internal state. Affected parameters include mcp_headers, conversation_id, user_conversation, token, and return type annotations. Examples: TaskResultAggregator.task_status_message returns Optional[Message]; A2AAgentExecutor.__init__ accepts Optional[dict[str, dict[str, str]]]; retrieve_response and prepare_tools_for_responses_api signatures updated throughout. Internal state variables like _TASK_STORE and _CONTEXT_STORE changed to Optional[...] declarations.
Utility module type annotation standardization src/utils/types.py	Updated ShieldModerationResult dataclass fields from union-with-None to Optional[...]: message and shield_model now declared as Optional[str] instead of `str

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• LCORE-1051: Updated Optional type in LCORE sources #978: Modifies Optional typing for A2A context storage-related declarations (_context_store, A2AContextStore), directly overlapping with type standardization scope in this PR.
• LCORE-956: Bump llama stack to 0.3.0 #866: Implements similar Optional-based type annotation changes across the same endpoint modules and utility files (query, query_v2, streaming_query, streaming_query_v2, utils/types).
• LCORE-740: type hints for endpoint handlers #752: Converts union-with-None type hints to Optional and strengthens type annotations on endpoint handler signatures, following the same refactoring pattern.

Suggested reviewers

• umago
• are-ces
• manstis

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Title check	❓ Inconclusive	The title 'LCORE-1051: final updates' is vague and does not clearly describe the main changes; it references a ticket but uses generic language like 'final updates' that obscures the actual nature of the changeset.	Provide a more descriptive title that summarizes the primary change, such as 'Standardize type annotations to use Optional across endpoints' or similar, to better communicate the purpose of this refactoring.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (3)

src/app/endpoints/query.py (3)

158-164: Inconsistent with coding guidelines for return types and local variables.

Per the coding guidelines, Optional[Type] should be used specifically for optional parameters, while return types and local variables should use the modern union syntax (Type | None). Line 159's parameter annotation is correct, but lines 161, 163, and 164 should use str | None.

♻️ Suggested change to align with guidelines

 def evaluate_model_hints(
     user_conversation: Optional[UserConversation],
     query_request: QueryRequest,
-) -> tuple[Optional[str], Optional[str]]:
+) -> tuple[str | None, str | None]:
     """Evaluate model hints from user conversation."""
-    model_id: Optional[str] = query_request.model
-    provider_id: Optional[str] = query_request.provider
+    model_id: str | None = query_request.model
+    provider_id: str | None = query_request.provider

Based on coding guidelines: "Use union types with modern syntax" and "Use Optional[Type] for optional parameters".

---

274-274: Local variable should use modern union syntax.

Per the coding guidelines, Optional[Type] is for optional function parameters. Local variables should use the modern union syntax.

♻️ Suggested change

-    user_conversation: Optional[UserConversation] = None
+    user_conversation: UserConversation | None = None

Based on coding guidelines: "Use union types with modern syntax: str | int instead of Union[str, int]".

---

860-874: Return type should use modern union syntax.

Per the coding guidelines, Optional[Type] is specifically for optional function parameters. Return types should use the modern union syntax Type | None.

♻️ Suggested change

 def get_rag_toolgroups(
     vector_db_ids: list[str],
-) -> Optional[list[Toolgroup]]:
+) -> list[Toolgroup] | None:
     """
     Return a list of RAG Tool groups if the given vector DB list is not empty.

     Generate a list containing a RAG knowledge search toolgroup if
     vector database IDs are provided.

     Parameters:
         vector_db_ids (list[str]): List of vector database identifiers to include in the toolgroup.

     Returns:
-        Optional[list[Toolgroup]]: A list with a single RAG toolgroup if
+        list[Toolgroup] | None: A list with a single RAG toolgroup if
         vector_db_ids is non-empty; otherwise, None.
     """

Based on coding guidelines: "Use union types with modern syntax: str | int instead of Union[str, int]".

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 06686ba and 1b7ddb1.

📒 Files selected for processing (6)

• src/app/endpoints/a2a.py
• src/app/endpoints/query.py
• src/app/endpoints/query_v2.py
• src/app/endpoints/streaming_query.py
• src/app/endpoints/streaming_query_v2.py
• src/utils/types.py

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Use FastAPI dependencies: from fastapi import APIRouter, HTTPException, Request, status, Depends
Use Llama Stack imports: from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules start with descriptive docstrings explaining purpose
Use logger = logging.getLogger(__name__) pattern for module logging
Type aliases defined at module level for clarity
All functions require docstrings with brief descriptions
Complete type annotations for function parameters and return types, using typing_extensions.Self for model validators
Use union types with modern syntax: str | int instead of Union[str, int]
Use Optional[Type] for optional parameters
Use snake_case with descriptive, action-oriented function names (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead
Use async def for I/O operations and external API calls
Handle APIConnectionError from Llama Stack in error handling
Use logger.debug() for detailed diagnostic information
Use logger.info() for general information about program execution
Use logger.warning() for unexpected events or potential problems
Use logger.error() for serious problems that prevented function execution
All classes require descriptive docstrings explaining purpose
Use PascalCase for class names with descriptive names and standard suffixes: Configuration, Error/Exception, Resolver, Interface
Use ABC for abstract base classes with @abstractmethod decorators
Complete type annotations for all class attributes
Follow Google Python docstring conventions (https://google.github.io/styleguide/pyguide.html) with sections: Args, Returns, Raises, Attributes
Run uv run make format to auto-format code with black and ruff before completion
Run uv run make verify to run all linters (black, pyl...

Files:

• src/utils/types.py
• src/app/endpoints/query.py
• src/app/endpoints/streaming_query_v2.py
• src/app/endpoints/streaming_query.py
• src/app/endpoints/a2a.py
• src/app/endpoints/query_v2.py

src/app/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Use FastAPI HTTPException with appropriate status codes for API endpoint error handling

Files:

• src/app/endpoints/query.py
• src/app/endpoints/streaming_query_v2.py
• src/app/endpoints/streaming_query.py
• src/app/endpoints/a2a.py
• src/app/endpoints/query_v2.py

🧠 Learnings (2)

📓 Common learnings

Learnt from: CR
Repo: lightspeed-core/lightspeed-stack PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-11T16:30:41.784Z
Learning: Applies to **/*.py : Use union types with modern syntax: `str | int` instead of `Union[str, int]`

Learnt from: CR
Repo: lightspeed-core/lightspeed-stack PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-11T16:30:41.784Z
Learning: Applies to **/*.py : Use `Optional[Type]` for optional parameters

Learnt from: CR
Repo: lightspeed-core/lightspeed-stack PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-11T16:30:41.784Z
Learning: Applies to **/*.py : Complete type annotations for function parameters and return types, using `typing_extensions.Self` for model validators

📚 Learning: 2026-01-11T16:30:41.784Z

Learnt from: CR
Repo: lightspeed-core/lightspeed-stack PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-01-11T16:30:41.784Z
Learning: Applies to **/*.py : Use FastAPI dependencies: `from fastapi import APIRouter, HTTPException, Request, status, Depends`

Applied to files:

• src/app/endpoints/streaming_query.py
• src/app/endpoints/a2a.py

🧬 Code graph analysis (4)

src/app/endpoints/query.py (1)

src/models/requests.py (1)

• QueryRequest (73-276)

src/app/endpoints/streaming_query.py (1)

src/models/database/conversations.py (1)

• UserConversation (11-38)

src/app/endpoints/a2a.py (1)

src/a2a_storage/context_store.py (1)

• A2AContextStore (7-58)

src/app/endpoints/query_v2.py (1)

src/utils/types.py (2)

• ToolCallSummary (111-119)
• ToolResultSummary (122-133)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: build-pr
• GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
• GitHub Check: E2E: server mode / ci
• GitHub Check: E2E: library mode / ci

🔇 Additional comments (29)

src/utils/types.py (1)

107-108: LGTM!

The type annotation changes from str | None to Optional[str] align with the coding guidelines for optional parameters. The Pydantic model fields are correctly annotated.

src/app/endpoints/streaming_query_v2.py (4)

4-4: LGTM!

Import updated to include Optional for the standardized type annotations used throughout the file.

---

141-141: LGTM!

Type annotation correctly updated to Optional[Any] for the nullable internal state variable.

---

375-376: LGTM!

Function parameter type annotation correctly updated to Optional[dict[str, dict[str, str]]] per coding guidelines.

---

474-475: LGTM!

Optional parameter type annotation correctly updated to Optional[str].

src/app/endpoints/a2a.py (6)

8-8: LGTM!

Import updated to include Optional for the standardized type annotations.

---

68-69: LGTM!

Module-level state variables correctly typed as Optional[TaskStore] and Optional[A2AContextStore] for the lazy initialization pattern.

---

123-123: LGTM!

Instance attribute correctly typed as Optional[Message] for the nullable state.

---

172-172: LGTM!

Property return type correctly updated to Optional[Message].

---

188-188: LGTM!

Constructor parameter type annotation correctly updated to Optional[dict[str, dict[str, str]]].

---

416-417: LGTM!

Method parameter type annotation correctly updated to Optional[str] for the nullable conversation ID.

src/app/endpoints/streaming_query.py (6)

10-18: LGTM!

Import block reformatted to include Optional for the standardized type annotations used throughout the file.

---

242-243: LGTM!

Function parameter type annotation correctly updated to Optional[str] for the nullable conversation ID.

---

395-396: LGTM!

Function parameter type annotation correctly updated to Optional[str].

---

745-745: LGTM!

Internal state variable correctly typed as Optional[Any].

---

861-861: LGTM!

Local variable correctly typed as Optional[UserConversation] for the nullable conversation lookup result.

---

1012-1013: LGTM!

Function parameter type annotation correctly updated to Optional[dict[str, dict[str, str]]].

src/app/endpoints/query_v2.py (10)

7-7: LGTM!

Import updated to include Optional for the standardized type annotations.

---

77-77: LGTM!

Return type annotation correctly updated to use Optional[ToolCallSummary] and Optional[ToolResultSummary] for the nullable tuple elements.

---

113-113: LGTM!

Local variable correctly typed as Optional[Any] for the nullable response payload.

---

297-298: LGTM!

Function parameter type annotation correctly updated to Optional[dict[str, dict[str, str]]].

---

508-508: LGTM!

Set type annotation correctly uses Optional[str] for nullable tuple elements in the deduplication set.

---

538-539: LGTM!

Comment and code correctly reference Optional[AnyUrl] for the nullable URL type.

---

695-706: LGTM!

Function return type and docstring correctly updated to Optional[list[dict[str, Any]]].

---

720-721: LGTM!

Function parameters correctly typed as Optional[str] and Optional[dict[str, dict[str, str]]].

---

765-766: LGTM!

Function parameter and return type correctly updated to use Optional typing.

---

781-782: LGTM!

Docstring correctly updated to reflect the Optional[list[dict[str, Any]]] return type.

src/app/endpoints/query.py (2)

485-487: LGTM!

Using Optional[str] for these optional function parameters aligns with the coding guideline: "Use Optional[Type] for optional parameters".

---

666-666: LGTM!

Using Optional[dict[str, dict[str, str]]] for this optional parameter with default None is correct per the coding guidelines.