promptbuilder 0.4.35__py3-none-any.whl → 0.4.37__py3-none-any.whl

promptbuilder/llm_client/base_client.py

@@ -319,34 +319,47 @@ class BaseLLMClient(ABC, utils.InheritDecoratorsMixin):
             max_tokens = self.default_max_tokens
         
         stream_messages = []
-
         total_count = 0
-        for response in self._create_stream(
-            messages=messages,
-            thinking_config=thinking_config,
-            system_message=system_message,
-            max_tokens=max_tokens if not autocomplete else None,
-        ):
-            yield response
+        response: Response | None = None
+
+        # Factory to (re)create the underlying provider stream using current accumulated state
+        def _stream_factory():
+            nonlocal response, total_count
+            tries = 3
+            while tries > 0:
+                try:
+                    iter = self._create_stream(
+                        messages=messages + stream_messages,
+                        thinking_config=thinking_config,
+                        system_message=system_message,
+                        max_tokens=max_tokens if not autocomplete else None,
+                    )
+                    for response in iter:
+                        yield response
+                    break
+                except Exception as e:
+                    tries -= 1
+                    if tries == 0:
+                        raise
+                    logger.warning(f"Stream generation error: {e}, retrying...")
+
+        # Use retry to iterate through the stream; on exception previously yielded parts
+        # are already merged into stream_messages so resumed attempts continue generation.
+        for response in _stream_factory():
             BaseLLMClient._append_generated_part(stream_messages, response)
-            finish_reason = response.candidates[0].finish_reason.value if response.candidates and response.candidates[0].finish_reason else None
             total_count += BaseLLMClient._response_out_tokens(response)
-            if finish_reason:
-                if autocomplete:
-                    while response.candidates and finish_reason == FinishReason.MAX_TOKENS.value:
-                        for response in self._create_stream(
-                            messages=messages,
-                            thinking_config=thinking_config,
-                            system_message=system_message,
-                            max_tokens=max_tokens if not autocomplete else None,
-                        ):
-                            yield response
-                            BaseLLMClient._append_generated_part(stream_messages, response)
-                            finish_reason = response.candidates[0].finish_reason.value if response.candidates and response.candidates[0].finish_reason else None
-                            total_count += BaseLLMClient._response_out_tokens(response)
-                        if max_tokens is not None and total_count >= max_tokens:
-                            break
-    
+            yield response
+        finish_reason = response.candidates[0].finish_reason.value if response and response.candidates and response.candidates[0].finish_reason else None
+        if finish_reason and autocomplete:
+            while response.candidates and finish_reason == FinishReason.MAX_TOKENS.value:
+                for response in _stream_factory():
+                    BaseLLMClient._append_generated_part(stream_messages, response)
+                    total_count += BaseLLMClient._response_out_tokens(response)
+                    yield response
+                finish_reason = response.candidates[0].finish_reason.value if response and response.candidates and response.candidates[0].finish_reason else None
+                if max_tokens is not None and total_count >= max_tokens:
+                    break
+
     @overload
     def from_text(
         self,
@@ -673,33 +686,44 @@ class BaseLLMClientAsync(ABC, utils.InheritDecoratorsMixin):
             max_tokens = self.default_max_tokens
         
         total_count = 0
-        stream_iter = await self._create_stream(
-            messages=messages,
-            thinking_config=thinking_config,
-            system_message=system_message,
-            max_tokens=max_tokens if not autocomplete else None,
-        )
+        response = None
+
+        async def _stream_factory():
+            nonlocal response, total_count
+            tries = 3
+            while tries > 0:
+                try:
+                    iter = await self._create_stream(
+                        messages=messages,
+                        thinking_config=thinking_config,
+                        system_message=system_message,
+                        max_tokens=max_tokens if not autocomplete else None,
+                    )
+
+                    async for response in iter:
+                        BaseLLMClient._append_generated_part(messages, response)
+                        total_count += BaseLLMClient._response_out_tokens(response)
+                        yield response
+                    break
+                except Exception as e:
+                    tries -= 1
+                    if tries <= 0:
+                        raise
+                    logger.warning(f"Stream generation error: {e}, retrying...")
+        
+        stream_iter = _stream_factory()
         async for response in stream_iter:
             yield response
-            BaseLLMClient._append_generated_part(messages, response)
-            finish_reason = response.candidates[0].finish_reason.value if response.candidates and response.candidates[0].finish_reason else None
-            total_count += BaseLLMClient._response_out_tokens(response)
-            if finish_reason:
-                if autocomplete:
-                    while response.candidates and finish_reason == FinishReason.MAX_TOKENS.value:
-                        stream_iter = await self._create_stream(
-                            messages=messages,
-                            thinking_config=thinking_config,
-                            system_message=system_message,
-                            max_tokens=max_tokens if not autocomplete else None,
-                        )
-                        async for response in stream_iter:
-                            yield response
-                            BaseLLMClient._append_generated_part(messages, response)
-                            finish_reason = response.candidates[0].finish_reason.value if response.candidates and response.candidates[0].finish_reason else None
-                            total_count += BaseLLMClient._response_out_tokens(response)
-                        if max_tokens is not None and total_count >= max_tokens:
-                            break
+
+        finish_reason = response.candidates[0].finish_reason.value if response and response.candidates and response.candidates[0].finish_reason else None
+        if finish_reason and autocomplete:
+            while response.candidates and finish_reason == FinishReason.MAX_TOKENS.value:
+                stream_iter = _stream_factory()
+                async for response in stream_iter:
+                    yield response
+                finish_reason = response.candidates[0].finish_reason.value if response.candidates and response.candidates[0].finish_reason else None
+                if max_tokens is not None and total_count >= max_tokens:
+                    break
 
     @overload
     async def from_text(

promptbuilder/llm_client/logfire_decorators.py

@@ -54,6 +54,21 @@ def extract_response_data(response: Response) -> dict[str, Any]:
     return response_data
 
 
+def record(span: logfire.LogfireSpan, duration: float, response: Response):
+    span.set_attribute("duration", duration)
+
+    span.set_attribute("response_data", extract_response_data(response))
+    span.set_attribute("candidates", response.candidates)
+    span.set_attribute("parsed", response.parsed)
+    span.set_attribute("response_text", response.text)
+    if response.usage_metadata is not None:
+        span.set_attribute("usage_metadata.cached_content_token_count", response.usage_metadata.cached_content_token_count)
+        span.set_attribute("usage_metadata.candidates_token_count", response.usage_metadata.candidates_token_count)
+        span.set_attribute("usage_metadata.thoughts_token_count", response.usage_metadata.thoughts_token_count)
+        span.set_attribute("usage_metadata.prompt_token_count", response.usage_metadata.prompt_token_count)
+        span.set_attribute("usage_metadata.total_token_count", response.usage_metadata.total_token_count)
+
+
 @inherited_decorator
 def create(class_method: Callable[P, Response]) -> Callable[P, Response]:
     """
@@ -69,17 +84,7 @@ def create(class_method: Callable[P, Response]) -> Callable[P, Response]:
         with logfire_llm.span(f"Create with {span_data["full_model_name"]}", **span_data) as span:
             start_time = time.time()
             response = class_method(self, *args, **kwargs)
-            span.set_attribute("duration", time.time() - start_time)
-            
-            span.set_attribute("response_data", extract_response_data(response))
-            span.set_attribute("candidates", response.candidates)
-            span.set_attribute("parsed", response.parsed)
-            span.set_attribute("response_text", response.text)
-            if response.usage_metadata is not None:
-                span.set_attribute("usage_metadata.cached_content_token_count", response.usage_metadata.cached_content_token_count)
-                span.set_attribute("usage_metadata.candidates_token_count", response.usage_metadata.candidates_token_count)
-                span.set_attribute("usage_metadata.prompt_token_count", response.usage_metadata.prompt_token_count)
-                span.set_attribute("usage_metadata.total_token_count", response.usage_metadata.total_token_count)
+            record(span, time.time() - start_time, response)
             
             return response
     
@@ -101,17 +106,7 @@ def create_async(class_method: Callable[P, Awaitable[Response]]) -> Callable[P,
         with logfire_llm.span(f"Async create with {span_data["full_model_name"]}", **span_data) as span:
             start_time = time.time()
             response = await class_method(self, *args, **kwargs)
-            span.set_attribute("duration", time.time() - start_time)
-            
-            span.set_attribute("response_data", extract_response_data(response))
-            span.set_attribute("candidates", response.candidates)
-            span.set_attribute("parsed", response.parsed)
-            span.set_attribute("response_text", response.text)
-            if response.usage_metadata is not None:
-                span.set_attribute("usage_metadata.cached_content_token_count", response.usage_metadata.cached_content_token_count)
-                span.set_attribute("usage_metadata.candidates_token_count", response.usage_metadata.candidates_token_count)
-                span.set_attribute("usage_metadata.prompt_token_count", response.usage_metadata.prompt_token_count)
-                span.set_attribute("usage_metadata.total_token_count", response.usage_metadata.total_token_count)
+            record(span, time.time() - start_time, response)
             
             return response
             
@@ -150,6 +145,7 @@ def record_streaming(span: logfire.LogfireSpan):
         span.set_attribute("response_text", stream_state.get_response_data()["message"]["content"])
         span.set_attribute("usage_metadata.cached_content_token_count", stream_state.last_usage_data.cached_content_token_count)
         span.set_attribute("usage_metadata.candidates_token_count", stream_state.last_usage_data.candidates_token_count)
+        span.set_attribute("usage_metadata.thoughts_token_count", stream_state.last_usage_data.thoughts_token_count)
         span.set_attribute("usage_metadata.prompt_token_count", stream_state.last_usage_data.prompt_token_count)
         span.set_attribute("usage_metadata.total_token_count", stream_state.last_usage_data.total_token_count)
 

{promptbuilder-0.4.35.dist-info → promptbuilder-0.4.37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: promptbuilder
-Version: 0.4.35
+Version: 0.4.37
 Summary: Library for building prompts for LLMs
 Home-page: https://github.com/kapulkin/promptbuilder
 Author: Kapulkin Stanislav

{promptbuilder-0.4.35.dist-info → promptbuilder-0.4.37.dist-info}/RECORD

@@ -9,20 +9,20 @@ promptbuilder/agent/utils.py,sha256=vTkphKw04v_QDIJtoB2JKK0RGY6iI1t_0LbmuStunzI,
 promptbuilder/llm_client/__init__.py,sha256=wJ33cnRtZX_YPsbcGxEu3SEZMOhPX7-fHI59MEPUe7I,517
 promptbuilder/llm_client/aisuite_client.py,sha256=8inY3UoH8o9yEOvRYP6a_8pjGQK0W_f9eV8MmHzpKTU,15641
 promptbuilder/llm_client/anthropic_client.py,sha256=YR1Pc4Fj0WpG7qcQnPLkQMzDsvA7SMvWomFR1oCzMsk,28328
-promptbuilder/llm_client/base_client.py,sha256=tpT_XlSZOW0PAB1OkIJQsDc5OUz0MdjB6GAfLjNW3Jk,34147
+promptbuilder/llm_client/base_client.py,sha256=RZ7AF5b0Z8Y8WACcyhSRBe4OJFjkhsDaCkiiUYkvYIM,34852
 promptbuilder/llm_client/bedrock_client.py,sha256=IQt7Sv_Wt6mg5-bhuyr-Nwjx5Nxk2S8rKEVkfwvWqE0,28183
 promptbuilder/llm_client/config.py,sha256=exQEm35wp7lK5SfXNpN5H9VZEb2LVa4pyZ-cxGt1U-U,1124
 promptbuilder/llm_client/exceptions.py,sha256=t-X7r_a8B1jNu8eEavde1jXu5dz97yV3IG4YHOtgh0Y,4836
 promptbuilder/llm_client/google_client.py,sha256=xXMWRf0TRuQCPmYfo2fnPRuoUnSiPNUOGU-GfHnd6uo,12234
 promptbuilder/llm_client/litellm_client.py,sha256=faQZ5k0tbQa0Z6t43OE05Ky1eRXquy9cVgMGpZCc2Uk,36416
-promptbuilder/llm_client/logfire_decorators.py,sha256=pEPm3nTfqkg9kZr3s5ORs3fMex_z4Ce8CAt301m3OIk,9697
+promptbuilder/llm_client/logfire_decorators.py,sha256=IoZf3PY1TFZlazY-GuAICF2GAJ8qtMpkTvBU3RPw9TQ,9126
 promptbuilder/llm_client/main.py,sha256=2Q7J5FwivX2YwvptzoSEtCfvfcI9p5HC55D3mMb2se4,8243
 promptbuilder/llm_client/openai_client.py,sha256=ZgI22-j4_B6OmvV10DO1lk4yN63hNYAaHx_K6E3pTSA,25276
 promptbuilder/llm_client/types.py,sha256=fnkSMFjK9ViaRQsD6LILpLz8R2_E1TI9efjy8VNO0RQ,8139
 promptbuilder/llm_client/utils.py,sha256=79lvSppjrrItHB5MIozbp_5Oq7TsOK4Qzt9Ae3XMLFw,7624
 promptbuilder/llm_client/vertex_client.py,sha256=OgbmRIYQXeK2kAh27ijhXuTzeg8APP38IyJ9WJqvDkY,15405
-promptbuilder-0.4.35.dist-info/licenses/LICENSE,sha256=fqXmInzgsvEOIaKSBgcrwKyYCGYF0MKErJ0YivtODcc,1096
-promptbuilder-0.4.35.dist-info/METADATA,sha256=I7gvoYDCgoHnFUB180N5LjBZlWPFg5u-rEn48oCfJv8,3799
-promptbuilder-0.4.35.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-promptbuilder-0.4.35.dist-info/top_level.txt,sha256=UBVcYn4UgrPy3O3fmmnPEU_kieuplBMgheetIMei4EI,14
-promptbuilder-0.4.35.dist-info/RECORD,,
+promptbuilder-0.4.37.dist-info/licenses/LICENSE,sha256=fqXmInzgsvEOIaKSBgcrwKyYCGYF0MKErJ0YivtODcc,1096
+promptbuilder-0.4.37.dist-info/METADATA,sha256=Qjz3KpGxhDKBbjJQpkNXLQCXNMEmAn4orrnraB6qKt4,3799
+promptbuilder-0.4.37.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+promptbuilder-0.4.37.dist-info/top_level.txt,sha256=UBVcYn4UgrPy3O3fmmnPEU_kieuplBMgheetIMei4EI,14
+promptbuilder-0.4.37.dist-info/RECORD,,