add gemini

app.py

@@ -8,6 +8,7 @@ from langchain_core.output_parsers import JsonOutputParser
 from langchain_openai.chat_models import ChatOpenAI
 from langchain.schema import HumanMessage, SystemMessage, AIMessage
 from langchain_anthropic import ChatAnthropic, ChatAnthropicMessages
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
 import openai
 
@@ -20,6 +21,8 @@ from langchain_openai import OpenAIEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.vectorstores import VectorStoreRetriever
 
+import time 
+
 
 feature_text = "Brand: Duckly. \nProduct name: Duck runner pro. \nKey properties: t-shirt, for running, sweat wicking, for marathon, 100% cotton."
 
@@ -85,7 +88,8 @@ models = ["gpt-4-turbo",
           "gpt-3.5-turbo", 
           "claude-3-sonnet-20240229", 
           "claude-3-opus-20240229", 
-          "claude-3-5-sonnet-20240620"
+          "claude-3-5-sonnet-20240620",
+          "gemini-1.5-pro"
           #"llama3-70b-8192",
           ]
 
@@ -229,13 +233,22 @@ def get_language(struct_lang, copy_lang):
     return "English"
 
 
-def get_model(model_name):
+def post_process(text):
+    # text = text.strip().replace('<br>', '  ')
+    with open(f"postprocess_{time.time()}.txt", "w") as f:
+        f.write(text)
+    return text
+
+
+def get_model(model_name, temperature=0.0):
     if model_name.startswith("gpt"):
-        chat = ChatOpenAI(model=model_name, max_tokens=4096)
+        chat = ChatOpenAI(model=model_name, max_tokens=4096, temperature=temperature)
     elif model_name.startswith("claude"):
-        chat = ChatAnthropic(model_name=model_name, anthropic_api_key=os.environ["ANTHROPIC_API_KEY"], max_tokens_to_sample=4096)
+        chat = ChatAnthropic(model_name=model_name, anthropic_api_key=os.environ["ANTHROPIC_API_KEY"], max_tokens_to_sample=4096, temperature=temperature)
+    elif model_name.startswith("gemini"):
+        chat = ChatGoogleGenerativeAI(model=model_name, api_key=os.environ["GOOGLE_API_KEY"], temperature=temperature)
     else:
-        chat = ChatGroq(model_name=model_name, api_key=os.environ["GROQ_API_KEY"])
+        chat = ChatGroq(model_name=model_name, api_key=os.environ["GROQ_API_KEY"], temperature=temperature)
     return chat
     
 
@@ -309,7 +322,7 @@ def generate(*data):
     if glossary_upload is not None:
         glossary = build_glossary(glossary_upload)
 
-    chat = get_model(model)
+    chat = get_model(model, temperature=temperature)
 
     types, struct_languages = parse_structure(struct_ref)
     copy_languages = detect_language([struct_ref[2 * i + 1] for i in range(visible + 1)], model=chat) 
@@ -363,7 +376,7 @@ Do not hallucinate any information.
 Use creative language in each output, do not use the common ways of starting product descriptions.
 Avoid common phrases and cliches such as "Step into something", "Elevate something", "Discover something", "Unleash something", "Embrace something", and similar phrases.
 For each version, try to write in different style.
-Put each item in the list on a separate line. Use <br> tag to separate items in the list.
+Put each item in the list on a separate line.  
 Rate the quality of each version based on the following criteria:
     - how well it follows the reference structure.
     - how well it follows the tone of voice, rythm, cadence and style of the reference copy.
@@ -427,7 +440,7 @@ Do not hallucinate any information.
 Use creative language in each output, do not use the common ways of starting product descriptions.
 Avoid common phrases and cliches such as "Step into something", "Elevate something", "Discover something", "Unleash something", "Embrace something", and similar phrases.
 For each version, try to write in different style.
-Put each item in the list on a separate line. Use <br> tag to separate items in the list.
+Put each item in the list on a separate line.  
 Rate the quality of each version based on the following criteria:
     - how well it follows the tone of voice, rythm, cadence and style of the reference copy.
     - how faithful it describes the product features.
@@ -490,7 +503,7 @@ Do not hallucinate any information.
 Use creative language in each output, do not use the common ways of starting product descriptions.
 Avoid common phrases and cliches such as "Step into something", "Elevate something", "Discover something", "Unleash something", "Embrace something", and similar phrases.
 For each version, try to write in different style.
-Put each item in the list on a separate line. Use <br> tag to separate items in the list.
+Put each item in the list on a separate line.  
 Rate the quality of each version based on the following criteria:
     - how well it follows the reference structure.
     - how faithful it describes the product features. 
@@ -538,7 +551,7 @@ Excluded words: {excluded_words}"""),]
 
     descriptions = ""
     
-    response = chat.batch(batch, temperature=temperature)
+    response = chat.batch(batch)
     print(response)
 
     parser = JsonOutputParser()
@@ -549,7 +562,7 @@ Excluded words: {excluded_words}"""),]
         bestid = jr["best_version"]["id"]
         for d in jr["versions"]:
             if d["id"] == bestid:
-                bestd = d["content"] + (f"\n\nDebug info:\n\nScore: {d['score']}\n\nExplanation: {jr['best_version']['explanation']}" if debug else "")
+                bestd = post_process(d["content"]) + (f"\n\nDebug info:\n\nScore: {d['score']}\n\nExplanation: {jr['best_version']['explanation']}" if debug else "")
                 bests  = d["score"]
                 break
         # bests = 0

requirements.txt

@@ -7,4 +7,5 @@ langchain_groq
 bcrypt
 langchain_chroma
 langchainhub
-langchain_community
+langchain_community
+langchain-google-genai