Fixes for claude/gemini calls (#18)

Co-authored-by: Kate Yanchenko <kate@lamoom.com>

Author: Katsiarynka

flow_prompt/ai_models/claude/claude_model.py

@@ -83,16 +83,26 @@ def get_client(self, client_secrets: dict) -> anthropic.Anthropic:
         return anthropic.Anthropic(api_key=client_secrets.get('api_key'))
 
 
+    def uny_all_messages_with_same_role(self, messages: t.List[dict]) -> t.List[dict]:
+        result = []
+        last_role = None
+        for message in messages:
+            if last_role != message.get("role"):
+                result.append(message)
+                last_role = message.get("role")
+            else:
+                result[-1]["content"] += message.get("content")
+        return result
+
+
     def call(self, messages: t.List[dict], max_tokens: int, client_secrets: dict = {}, **kwargs) -> AIResponse:
         common_args = get_common_args(max_tokens)
         kwargs = {
-            **{
-                "messages": messages,
-            },
             **common_args,
             **self.get_params(),
             **kwargs,
         }
+        messages = self.uny_all_messages_with_same_role(messages)
 
         logger.debug(
             f"Calling {messages} with max_tokens {max_tokens} and kwargs {kwargs}"

flow_prompt/ai_models/gemini/gemini_model.py

@@ -91,25 +91,24 @@ def call(self, messages: t.List[dict], max_tokens: int, client_secrets: dict = {
         stream_params = kwargs.get("stream_params")
 
         # Parse only prompt content due to gemini call specifics
-        prompts = [obj["content"] for obj in messages if obj["role"] == "user"]
+        prompt = '\n\n'.join([obj["content"] for obj in messages])
 
         content = ""
 
         try:
             if not kwargs.get('stream'):
-                for prompt in prompts:
-                    response = self.model.generate_content(prompt, stream=False)
-                    content = response.text
+                response = self.model.generate_content(prompt, stream=False)
+                content = response.text
             else:
+                response = self.model.generate_content(prompt, stream=True)
                 idx = 0
-                for prompt in prompts:
+                for chunk in response:
                     if idx % 5 == 0:
+                        idx = 0
                         if not check_connection(**stream_params):
                             raise ConnectionLostError("Connection was lost!")
-                    response = self.model.generate_content(prompt, stream=True)
-                    for chunk in response:
-                        stream_function(chunk.text, **stream_params)
-                        content += chunk.text
+                    stream_function(chunk.text, **stream_params)
+                    content += chunk.text
                     idx += 1
 
             return GeminiAIResponse(

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flow-prompt"
-version = "0.1.14a17"
+version = "0.1.20"
 description = ""
 authors = ["Flow-prompt Engineering Team <engineering@flow-prompt.com>"]
 readme = "README.md"