test

langchain_pipeline.py

@@ -1,4 +1,5 @@
 import os
+import requests
 from pdfminer import high_level
 from langchain_astradb import AstraDBVectorStore
 from langchain_core.prompts import PromptTemplate
@@ -6,18 +7,60 @@ from langchain_openai import OpenAIEmbeddings
 from langchain_anthropic import ChatAnthropic
 from langchain_google_genai import ChatGoogleGenerativeAI
 import base64
-from google.cloud import aiplatform
+import json
+
+# Load Google Cloud credentials from the file
+credentials_path = "/mnt/data/Composure IAM Admin.json"
+with open(credentials_path) as f:
+    credentials = json.load(f)
+
+def get_access_token():
+    # Request an access token using the service account credentials
+    auth_url = 'https://oauth2.googleapis.com/token'
+    auth_data = {
+        'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',
+        'assertion': create_jwt_assertion(credentials)
+    }
+    response = requests.post(auth_url, data=auth_data)
+    response.raise_for_status()
+    return response.json()['access_token']
+
+def create_jwt_assertion(credentials):
+    # Create a JWT assertion to use for requesting an access token
+    header = {
+        'alg': 'RS256',
+        'typ': 'JWT'
+    }
+    now = int(time.time())
+    claim_set = {
+        'iss': credentials['client_email'],
+        'scope': 'https://www.googleapis.com/auth/cloud-platform',
+        'aud': 'https://oauth2.googleapis.com/token',
+        'exp': now + 3600,
+        'iat': now
+    }
+    jwt_message = json.dumps(header) + '.' + json.dumps(claim_set)
+    signed_jwt = jwt.encode(jwt_message, credentials['private_key'], algorithm='RS256')
+    return signed_jwt
 
 def generate(document_parts, prompt_text):
-    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")
+    access_token = get_access_token()
+    model_endpoint = f'https://us-central1-aiplatform.googleapis.com/v1/projects/{credentials["project_id"]}/locations/us-central1/endpoints/gemini-1.5-pro-001:predict'
+    headers = {
+        'Authorization': f'Bearer {access_token}',
+        'Content-Type': 'application/json'
+    }
     instances = [{"content": part} for part in document_parts]
-    
-    response = client.predict(name=model, instances=instances, parameters={"temperature": 1.0})
-    
+    data = {
+        'instances': instances,
+        'parameters': {"temperature": 1.0}
+    }
+    response = requests.post(model_endpoint, headers=headers, json=data)
+    response.raise_for_status()
+    predictions = response.json()['predictions']
+
     response_text = ""
-    for prediction in response.predictions:
+    for prediction in predictions:
         response_text += prediction.get("content", "")
     
     return response_text