promptbuilder 0.4.35__tar.gz → 0.4.36__tar.gz

{promptbuilder-0.4.35/promptbuilder.egg-info → promptbuilder-0.4.36}/PKG-INFO

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

{promptbuilder-0.4.35 → promptbuilder-0.4.36}/promptbuilder/llm_client/base_client.py

@@ -321,32 +321,32 @@ class BaseLLMClient(ABC, utils.InheritDecoratorsMixin):
         stream_messages = []
 
         total_count = 0
+        response = None
         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 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 self._create_stream(
+                    messages=messages,
+                    thinking_config=thinking_config,
+                    system_message=system_message,
+                    max_tokens=max_tokens if not autocomplete else None,
+                ):
+                    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.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,
@@ -679,27 +679,28 @@ class BaseLLMClientAsync(ABC, utils.InheritDecoratorsMixin):
             system_message=system_message,
             max_tokens=max_tokens if not autocomplete else None,
         )
+        response = 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 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
+            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:
+                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)
+                    total_count += BaseLLMClient._response_out_tokens(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-0.4.35 → promptbuilder-0.4.36}/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 → promptbuilder-0.4.36/promptbuilder.egg-info}/PKG-INFO

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

{promptbuilder-0.4.35 → promptbuilder-0.4.36}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="promptbuilder",
-    version="0.4.35",
+    version="0.4.36",
     packages=find_packages(),
     install_requires=[
         "pydantic",