flo_ai/flo_ai/llm/rootflo_llm.py (4)

87-107: Lazy state storage and BaseLLM placeholder model/api_key look reasonable

Persisting constructor parameters (_base_url, _model_id, JWT fields, temperature, **kwargs) for later use and introducing _llm/_initialized flags align well with the new lazy-init flow. Calling BaseLLM.__init__ with placeholder model/api_key and then overwriting self.model and self.api_key in _ensure_initialized is a pragmatic way to preserve the parent’s setup without doing eager I/O in __init__.

One tiny nit: the comment “Will be set during initialization” above self.base_url/self.model_id is now slightly misleading since they’re set immediately, but this is non-blocking and can be tidied later if desired.

Also applies to: 108-115

---

117-169: Async config fetch via httpx.AsyncClient is correct; consider more specific exception types

The async config fetch logic is solid: correct URL construction, auth headers with JWT, a reasonable timeout, and explicit shape validation for data.llm_model and data.type. Switching from a sync client to httpx.AsyncClient should avoid the earlier deadlock if callers share the same base_url as flo-ai.

One improvement you might want to consider later is raising a more specific custom exception (e.g., RootFloLLMConfigError) instead of generic Exception so callers can distinguish config/bootstrap failures from downstream LLM errors more easily.

---

171-264: _ensure_initialized flow looks correct; minor robustness tweaks possible

The lazy-init sequence here is well structured:

• Fast-path and double-checked locking avoid redundant work under concurrency.
• JWT generation or reuse of a provided access_token is clean; api_token is then propagated consistently to the config fetch and downstream SDK wrappers.
• Config response is validated, llm_type normalized via LLMProvider, and misconfigurations surface as explicit ValueErrors.
• The full proxy inference URL and X-Rootflo-Key headers are built once and passed into the appropriate wrapper (OpenAI, Anthropic, Gemini), and _initialized is only set after _llm is constructed.

Two small robustness ideas:

1. Surface provider/type in errors – in the else branch at lines 259-260, you might include both the raw llm_type and enum value in the error to aid debugging, though this is non-essential.
2. Optionally cache config fields – if you foresee needing raw llm_type/llm_model from the API for introspection, you could cache self._llm_model_from_config / self._llm_type_from_config, but it’s not required for correctness.

Overall, this method cleanly centralizes all initialization concerns and matches the PR’s goal.

---

289-303: Guard synchronous helper methods against uninitialized _llm

With the shift to lazy initialization, it’s now possible (at least in theory) for callers to access get_message_content, format_tool_for_llm, format_tools_for_llm, or format_image_in_message before generate/stream has been called. In that case, self._llm would still be None, and these methods would raise an AttributeError instead of a clearer message.

To make failure modes more predictable, consider adding a simple guard:

    def get_message_content(self, response: Any) -> str:
        """Extract message content from response"""
-        return self._llm.get_message_content(response)
+        if self._llm is None:
+            raise RuntimeError(
+                "RootFloLLM is not initialized. Call 'generate' or 'stream' "
+                "at least once before using helper methods."
+            )
+        return self._llm.get_message_content(response)

and similarly for the other helper methods, or at least document that generate/stream must be called first. This isn’t strictly required if all current call sites follow that pattern, but it will make the class more defensive and easier to use correctly.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flo_ai/flo_ai/llm/rootflo_llm.py (2)

98-102: Potential event‑loop dependency when creating asyncio.Lock() in __init__

Instantiating asyncio.Lock() in a synchronous __init__ can raise RuntimeError (“no running event loop”) depending on the Python version and how/when RootFloLLM is constructed (e.g., at import time or app startup before an event loop is running). This would surface as an immediate construction failure.

Consider deferring lock creation until first use inside _ensure_initialized, which always runs in an async context, e.g.:

-        self._init_lock = asyncio.Lock()
+        self._init_lock = None

and in _ensure_initialized:

-        # Acquire lock for initialization
-        async with self._init_lock:
+        # Lazily create the lock in an async context
+        if self._init_lock is None:
+            self._init_lock = asyncio.Lock()
+
+        # Acquire lock for initialization
+        async with self._init_lock:

Please verify how RootFloLLM is constructed in your deployment (startup vs inside an async request handler) to confirm whether this is observable under your Python runtime.

---

171-264: Initialization guard and provider dispatch are correct, but helper methods can observe uninitialized _llm

The double‑checked locking in _ensure_initialized, JWT generation, config fetch, mapping type → LLMProvider, and instantiation of OpenAI/Anthropic/Gemini all look correct, and the fast‑path if self._initialized: return keeps steady‑state overhead low.

However, get_message_content, format_tool_for_llm, format_tools_for_llm, and format_image_in_message access self._llm without calling _ensure_initialized. If any caller uses these helpers before the first generate/stream (e.g., formats tools or images before sending the first request), self._llm will still be None and you’ll get an AttributeError.

Two options you might consider:

1. Defensive guard with clearer error (no interface change, just better failure mode):

    def get_message_content(self, response: Any) -> str:
        """Extract message content from response"""
-        return self._llm.get_message_content(response)
+        if not getattr(self, "_initialized", False) or self._llm is None:
+            raise RuntimeError(
+                "RootFloLLM is not initialized yet; call generate() or stream() first."
+            )
+        return self._llm.get_message_content(response)

and similarly for the other helper methods.

2. Contract clarification: If your current call‑sites always invoke generate/stream first, document that these helpers must only be used after initialization, and consider adding an explicit async initialize() method that some components can await upfront.

Please check existing usages (e.g., any code that calls format_tools_for_llm or format_image_in_message immediately after constructing the LLM) to make sure this change doesn’t introduce a latent NoneType access.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flo_ai/flo_ai/llm/rootflo_llm.py (1)

301-319: Guard helper methods against uninitialized _llm (runtime bug & prior feedback)

get_message_content now raises a clear RuntimeError when the LLM isn’t initialized, but the helper methods:

• format_tool_for_llm
• format_tools_for_llm
• format_image_in_message

still call self._llm.* directly. If any of these are used before generate() / stream() are awaited, they will raise an AttributeError: 'NoneType' object has no attribute ... instead of a clear, intentional error. In the previous eager-init design these helpers were usable immediately after construction; with lazy init this becomes a regression and matches an unresolved earlier review comment.

To align behavior and fail fast with a clear message, factor the guard into a small private helper and reuse it in all four methods, e.g.:

@@
-    def get_message_content(self, response: Any) -> str:
-        """Extract message content from response"""
-        if not getattr(self, '_initialized', False) or self._llm is None:
-            raise RuntimeError(
-                'RootFloLLM is not initialized yet; call generate() or stream() first.'
-            )
-        return self._llm.get_message_content(response)
+    def _require_initialized(self) -> None:
+        if not getattr(self, '_initialized', False) or self._llm is None:
+            raise RuntimeError(
+                'RootFloLLM is not initialized yet; call generate() or stream() first.'
+            )
+
+    def get_message_content(self, response: Any) -> str:
+        """Extract message content from response"""
+        self._require_initialized()
+        return self._llm.get_message_content(response)
@@
     def format_tool_for_llm(self, tool: 'Tool') -> Dict[str, Any]:
         """Format a tool for the specific LLM's API"""
-        return self._llm.format_tool_for_llm(tool)
+        self._require_initialized()
+        return self._llm.format_tool_for_llm(tool)
@@
     def format_tools_for_llm(self, tools: List['Tool']) -> List[Dict[str, Any]]:
         """Format a list of tools for the specific LLM's API"""
-        return self._llm.format_tools_for_llm(tools)
+        self._require_initialized()
+        return self._llm.format_tools_for_llm(tools)
@@
     def format_image_in_message(self, image: ImageMessageContent) -> str:
         """Format a image in the message"""
-        return self._llm.format_image_in_message(image)
+        self._require_initialized()
+        return self._llm.format_image_in_message(image)

This preserves lazy initialization semantics while ensuring all helper entrypoints either work or raise the same explicit RuntimeError instead of a confusing AttributeError.

flo_ai/flo_ai/llm/rootflo_llm.py (1)

16-23: Keep provider docs in sync with new VLLM support

Enum and wiring now support LLMProvider.VLLM, but the class docstring still says this proxies only “OpenAI, Gemini, and Anthropic”. Consider updating the docstring to include VLLM (and mention lazy config fetch) so callers don’t miss the new provider.

Also applies to: 27-29

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro