sendit

langchain_pipeline.py

@@ -6,23 +6,20 @@ from langchain_openai import OpenAIEmbeddings
 from langchain_anthropic import ChatAnthropic
 from langchain_google_genai import ChatGoogleGenerativeAI
 import base64
-import vertexai
-from vertexai.generative_models import GenerativeModel, Part, FinishReason
-import vertexai.preview.generative_models as generative_models
+from google.cloud import aiplatform
 
 def generate(document_parts, prompt_text):
-    vertexai.init(project="akroda", location="us-central1")
-    model = GenerativeModel("gemini-1.5-pro-001")
-    responses = model.generate_content(
-        document_parts,
-        generation_config=generation_config,
-        safety_settings=safety_settings,
-        stream=True,
-    )
-
+    aiplatform.init(project="akroda", location="us-central1")
+    client = aiplatform.PredictionServiceClient()
+    model = client.model_path(project="akroda", location="us-central1", model="gemini-1.5-pro-001")
+    instances = [{"content": part} for part in document_parts]
+    
+    response = client.predict(name=model, instances=instances, parameters={"temperature": 1.0})
+    
     response_text = ""
-    for response in responses:
-        response_text += response.text
+    for prediction in response.predictions:
+        response_text += prediction.get("content", "")
+    
     return response_text
 
 document1 = Part.from_data(
@@ -102,10 +99,10 @@ generation_config = {
 }
 
 safety_settings = {
-    generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-    generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+    "HARM_CATEGORY_HATE_SPEECH": "BLOCK_MEDIUM_AND_ABOVE",
+    "HARM_CATEGORY_DANGEROUS_CONTENT": "BLOCK_MEDIUM_AND_ABOVE",
+    "HARM_CATEGORY_SEXUALLY_EXPLICIT": "BLOCK_MEDIUM_AND_ABOVE",
+    "HARM_CATEGORY_HARASSMENT": "BLOCK_MEDIUM_AND_ABOVE",
 }
 
 def pipeline(file, model_name, balance_type, apsn_transactions, max_fees_per_day, min_overdrawn_fee, min_transaction_overdraft):