Update TextGen/router.py

TextGen/router.py

@@ -1,12 +1,14 @@
 from pydantic import BaseModel
 
-from .ConfigEnv import config
 from fastapi.middleware.cors import CORSMiddleware
 
-from langchain.llms import Clarifai
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
-
+from langchain_google_genai import (
+    ChatGoogleGenerativeAI,
+    HarmBlockThreshold,
+    HarmCategory,
+)
 from TextGen import app
 
 class Generate(BaseModel):
@@ -18,12 +20,12 @@ def generate_text(prompt: str):
     else:
         prompt = PromptTemplate(template=prompt, input_variables=['Prompt'])
 
-        llm = Clarifai(
-            pat = config.CLARIFAI_PAT,
-            user_id = config.USER_ID,
-            app_id = config.APP_ID, 
-            model_id = config.MODEL_ID,
-            model_version_id=config.MODEL_VERSION_ID,
+        # Initialize the LLM
+        llm = ChatGoogleGenerativeAI(
+            model="gemini-pro",
+            safety_settings={
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+            },
         )
 
         llmchain = LLMChain(