Update ChatCohere.py

ChatCohere.py

@@ -1,5 +1,4 @@
 import cohere
-from typing import List, Dict
 from dotenv import load_dotenv
 import os
 
@@ -8,19 +7,25 @@ load_dotenv()
 cohere_api_key = os.getenv("cohere_api_key")
 co = cohere.ClientV2(cohere_api_key)
 
-def chat_completion(message: str) -> str:
+def chat_completion(message: str, context: str) -> str:
     response = co.chat(
         model="command-r-plus-08-2024",
-        messages=[{"role": "system", "content": "You are a Pokemon expert. You know everything about their characteristics, abilities, and evolutions. You always reply with the most accurate information, you are polite and helpful. Your output cannot be larger than 1500 characters (spaces included)."}, {"role": "user", "content": message}],
+        messages=[{"role": "system", "content": "You are a Pokemon expert. You know everything about their characteristics, abilities, and evolutions. You always reply with the most accurate information, you are polite and helpful. Your output cannot be larger than 1500 characters (spaces included)."}, {"role": "user", "content": message}, {"role": "assistant", "content": context}],
     ) 
     return response.message.content[0].text
 
-def summarize(message, system_prompt="You are an helpful assistant whose job is to summarize the text you are given in less than 900 charachters (including spaces) in such a way that it would constitute an effective and engaging description of a pokemon card package"):
+def card_completion(message: str) -> str:
     response = co.chat(
         model="command-r-plus-08-2024",
-        messages=[
+        messages=[{"role": "system", "content": "You are an expert in Pokemon cards. You know everything about their characteristics, abilities, and evolutions. You always reply with the most accurate information, you are polite and helpful."}, {"role": "user", "content": message}],
+    ) 
+    return response.message.content[0].text
+
+def summarize(message, system_prompt="You are an helpful assistant whose job is to summarize the text you are given in less than 900 charachters (including spaces) in such a way that it would constitute an effective and engaging description of a pokemon card package"):
+        response = chat_completion(
+            message_history=[
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": message},
             ]
-    )
-    return response.message.content[0].text
+        )
+        return response