tanta

langchain_pipeline.py

@@ -9,47 +9,24 @@ from langchain_openai import OpenAIEmbeddings
 from langchain_anthropic import ChatAnthropic
 from langchain_google_genai import ChatGoogleGenerativeAI
 import base64
-import jwt  # Add PyJWT to your requirements
+from google.oauth2 import service_account
+from google.auth.transport.requests import Request
 
 # Load Google Cloud credentials from the environment variable
 credentials_json = os.getenv('GOOGLE_APPLICATION_CREDENTIALS_JSON')
 if credentials_json:
-    credentials = json.loads(credentials_json)
+    credentials_info = json.loads(credentials_json)
 else:
     raise ValueError("GOOGLE_APPLICATION_CREDENTIALS_JSON environment variable not set")
 
 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
+    credentials = service_account.Credentials.from_service_account_info(credentials_info)
+    credentials.refresh(Request())
+    return credentials.token
 
 def generate(document_parts, prompt_text):
     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'
+    model_endpoint = f'https://us-central1-aiplatform.googleapis.com/v1/projects/{credentials_info["project_id"]}/locations/us-central1/endpoints/gemini-1.5-pro-001:predict'
     headers = {
         'Authorization': f'Bearer {access_token}',
         'Content-Type': 'application/json'

requirements.txt

@@ -10,5 +10,4 @@ langchain-astradb
 langchain-openai
 langchain-anthropic
 langchain-google-genai
-docxcompose
-pyjwt
+docxcompose