chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,168 @@
|
||||
"""Cloud Function code to analyze a prospectus"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
|
||||
import functions_framework
|
||||
from google.cloud.alloydb.connector import Connector
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from langchain_google_vertexai import VertexAI
|
||||
import sqlalchemy
|
||||
|
||||
|
||||
# Triggered from a message on a Cloud Pub/Sub topic.
|
||||
@functions_framework.cloud_event
|
||||
def analyze_prospectus(cloud_event):
|
||||
"""Function to analyze prospectus"""
|
||||
# Print out the data from Pub/Sub, to prove that it worked
|
||||
ticker = base64.b64decode(cloud_event.data["message"]["data"])
|
||||
ticker = ticker.decode("utf-8")
|
||||
print(ticker)
|
||||
|
||||
# Environment Vars
|
||||
region = os.environ["REGION"]
|
||||
project_id = os.environ["PROJECT_ID"]
|
||||
|
||||
# AlloyDB Vars
|
||||
cluster = "alloydb-cluster"
|
||||
instance = "alloydb-instance"
|
||||
database = "ragdemos"
|
||||
table_name = "langchain_vector_store"
|
||||
user = "postgres"
|
||||
password = os.environ["ALLOYDB_PASSWORD"]
|
||||
|
||||
# Setup sync connector
|
||||
connector = Connector()
|
||||
|
||||
def getconn():
|
||||
conn = connector.connect(
|
||||
f"projects/{project_id}/locations/{region}/clusters/{cluster}/instances/{instance}",
|
||||
"pg8000",
|
||||
user=user,
|
||||
password=password,
|
||||
db=database,
|
||||
)
|
||||
return conn
|
||||
|
||||
# create connection pool
|
||||
pool = sqlalchemy.create_engine(
|
||||
"postgresql+pg8000://",
|
||||
creator=getconn,
|
||||
)
|
||||
|
||||
# Prep SQL statement
|
||||
sql = f"SELECT content FROM {table_name} WHERE ticker = '{ticker}' ORDER BY page, page_chunk"
|
||||
|
||||
# Prep model and template
|
||||
model = VertexAI(
|
||||
model_name="gemini-2.0-flash", max_output_tokens=1024, temperature=0.0
|
||||
)
|
||||
template = """
|
||||
<MISSION>
|
||||
You are an experienced financial analyst. Your mission is to create a detailed
|
||||
company financial overview for {ticker} using their latest prospectus. I will be
|
||||
sending you the prospectus a few chunks at a time. There are a total of
|
||||
{total_chunk_count} prospectus chunks, and I am sending you prospectus chunk numbers
|
||||
{first_chunk}-{last_chunk} as part of this request.
|
||||
</MISSION>
|
||||
|
||||
<TASK>
|
||||
Use the financial overview labeled <OVERVIEW> below, and use the additional details from
|
||||
the section labeled <ADDITIONAL_CONTEXT> below to improve the financial overview in the <OVERVIEW>.
|
||||
Respond using less than 4000 characters, including whitespace.
|
||||
</TASK>
|
||||
|
||||
<OVERVIEW>
|
||||
{previous_overview}
|
||||
</OVERVIEW>
|
||||
|
||||
<ADDITIONAL_CONTEXT>
|
||||
{chunk_text}
|
||||
</ADDITIONAL_CONTEXT>"""
|
||||
|
||||
prompt = PromptTemplate.from_template(template)
|
||||
|
||||
# Create overview of full document by iterating through chunks
|
||||
with pool.connect() as db_conn:
|
||||
# query database
|
||||
result = db_conn.execute(sqlalchemy.text(sql)).fetchall()
|
||||
|
||||
# commit transaction (SQLAlchemy v2.X.X is commit as you go)
|
||||
db_conn.commit()
|
||||
|
||||
# Iterate through results
|
||||
total_chunk_count = len(result)
|
||||
overview = ""
|
||||
chunk_text = ""
|
||||
first_chunk = 1
|
||||
last_chunk = 1
|
||||
|
||||
for i in range(len(result)):
|
||||
current_chunk = i + 1
|
||||
first_chunk = min(first_chunk, current_chunk)
|
||||
last_chunk = max(last_chunk, current_chunk)
|
||||
|
||||
# Add text to chunk_text until token window is full
|
||||
chunk_text = chunk_text + str(result[i].content) + " "
|
||||
if len(chunk_text) < 50000:
|
||||
continue
|
||||
|
||||
# Invoke the model
|
||||
print(
|
||||
f"Adding chunks {first_chunk} through {last_chunk} out of {total_chunk_count} to {ticker} overview..."
|
||||
)
|
||||
fmt_prompt = prompt.format(
|
||||
total_chunk_count=total_chunk_count,
|
||||
first_chunk=first_chunk,
|
||||
last_chunk=last_chunk,
|
||||
previous_overview=overview,
|
||||
chunk_text=chunk_text,
|
||||
ticker=ticker,
|
||||
)
|
||||
|
||||
overview = model.invoke(fmt_prompt)
|
||||
|
||||
# Reset first_chunk and chunk_text values
|
||||
first_chunk = current_chunk + 1
|
||||
chunk_text = ""
|
||||
|
||||
analysis = model.invoke(
|
||||
f"You are an experienced financial analyst. Write a financial analysis for ticker {ticker} that includes an Investment Rating (buy, sell, or hold), Investment Risk (high, medium, low), Target Investor (conservative, neutral, aggressive) and a two-paragraph analysis. Use the following company overview as context for the analysis: \n\n{overview}"
|
||||
)
|
||||
rating = model.invoke(
|
||||
f"Answering with only 1 word, classify ticker {ticker} as one of [BUY, SELL, HOLD] based on the following analysis: {analysis}"
|
||||
)
|
||||
rating = rating.strip()
|
||||
|
||||
insert_stmt = sqlalchemy.text(
|
||||
"INSERT INTO investments (id, ticker, etf, market, rating, overview, analysis) VALUES (:id, :ticker, :etf, :market, :rating, :overview, :analysis)"
|
||||
)
|
||||
|
||||
with pool.connect() as db_conn:
|
||||
max_id = db_conn.execute(
|
||||
sqlalchemy.text("SELECT MAX(id) FROM investments")
|
||||
).fetchall()
|
||||
new_id = max_id[0][0] + 1
|
||||
print(new_id)
|
||||
|
||||
# insert into database
|
||||
db_conn.execute(
|
||||
insert_stmt,
|
||||
parameters={
|
||||
"id": new_id,
|
||||
"ticker": ticker,
|
||||
"etf": False,
|
||||
"market": "US",
|
||||
"rating": rating,
|
||||
"overview": overview,
|
||||
"analysis": analysis,
|
||||
},
|
||||
)
|
||||
|
||||
# commit transaction (SQLAlchemy v2.X.X is commit as you go)
|
||||
db_conn.commit()
|
||||
print("Finished insert")
|
||||
|
||||
print("Closing database connection.")
|
||||
connector.close()
|
||||
print(f"Finished analyzing ticker {ticker}.")
|
||||
@@ -0,0 +1,5 @@
|
||||
functions-framework==3.*
|
||||
google-cloud-alloydb-connector[pg8000]==1.4.0
|
||||
langchain-core==0.3.31
|
||||
langchain-google-vertexai==2.0.15
|
||||
SQLAlchemy==2.0.34
|
||||
@@ -0,0 +1,300 @@
|
||||
"""Cloud Function code to process a pdf dropped in GCS"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import uuid
|
||||
|
||||
import functions_framework
|
||||
from google.api_core.client_options import ClientOptions
|
||||
from google.api_core.exceptions import InternalServerError, RetryError
|
||||
from google.cloud import documentai # type: ignore
|
||||
from google.cloud import pubsub_v1, storage
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_google_alloydb_pg import AlloyDBEngine, AlloyDBVectorStore, Column
|
||||
from langchain_google_vertexai import VertexAIEmbeddings
|
||||
|
||||
|
||||
# Source: https://cloud.google.com/document-ai/docs/samples/documentai-batch-process-document#documentai_batch_process_document-python
|
||||
def batch_process_documents(
|
||||
project_id: str,
|
||||
location: str,
|
||||
processor_id: str,
|
||||
gcs_output_uri: str,
|
||||
processor_version_id: str | None = None,
|
||||
gcs_input_uri: str | None = None,
|
||||
input_mime_type: str | None = None,
|
||||
gcs_input_prefix: str | None = None,
|
||||
field_mask: str | None = None,
|
||||
timeout: int = 400,
|
||||
) -> list[storage.Blob]:
|
||||
"""Function to batch process documents"""
|
||||
# You must set the `api_endpoint` if you use a location other than "us".
|
||||
opts = ClientOptions(api_endpoint=f"{location}-documentai.googleapis.com")
|
||||
|
||||
client = documentai.DocumentProcessorServiceClient(client_options=opts)
|
||||
|
||||
if gcs_input_uri:
|
||||
# Specify specific GCS URIs to process individual documents
|
||||
gcs_document = documentai.GcsDocument(
|
||||
gcs_uri=gcs_input_uri, mime_type=input_mime_type
|
||||
)
|
||||
# Load GCS Input URI into a List of document files
|
||||
gcs_documents = documentai.GcsDocuments(documents=[gcs_document])
|
||||
input_config = documentai.BatchDocumentsInputConfig(gcs_documents=gcs_documents)
|
||||
else:
|
||||
# Specify a GCS URI Prefix to process an entire directory
|
||||
gcs_prefix = documentai.GcsPrefix(gcs_uri_prefix=gcs_input_prefix)
|
||||
input_config = documentai.BatchDocumentsInputConfig(gcs_prefix=gcs_prefix)
|
||||
|
||||
# Cloud Storage URI for the Output Directory
|
||||
gcs_output_config = documentai.DocumentOutputConfig.GcsOutputConfig(
|
||||
gcs_uri=gcs_output_uri, field_mask=field_mask
|
||||
)
|
||||
|
||||
# Where to write results
|
||||
output_config = documentai.DocumentOutputConfig(gcs_output_config=gcs_output_config)
|
||||
|
||||
if processor_version_id:
|
||||
# The full resource name of the processor version, e.g.:
|
||||
# projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}
|
||||
name = client.processor_version_path(
|
||||
project_id, location, processor_id, processor_version_id
|
||||
)
|
||||
else:
|
||||
# The full resource name of the processor, e.g.:
|
||||
# projects/{project_id}/locations/{location}/processors/{processor_id}
|
||||
name = client.processor_path(project_id, location, processor_id)
|
||||
|
||||
request = documentai.BatchProcessRequest(
|
||||
name=name,
|
||||
input_documents=input_config,
|
||||
document_output_config=output_config,
|
||||
)
|
||||
|
||||
# BatchProcess returns a Long Running Operation (LRO)
|
||||
operation = client.batch_process_documents(request)
|
||||
|
||||
# Continually polls the operation until it is complete.
|
||||
# This could take some time for larger files
|
||||
# Format: projects/{project_id}/locations/{location}/operations/{operation_id}
|
||||
try:
|
||||
print(f"Waiting for operation {operation.operation.name} to complete...")
|
||||
operation.result(timeout=timeout)
|
||||
# Catch exception when operation doesn't finish before timeout
|
||||
except (RetryError, InternalServerError) as e:
|
||||
print(e.message)
|
||||
|
||||
# Once the operation is complete,
|
||||
# get output document information from operation metadata
|
||||
metadata = documentai.BatchProcessMetadata(operation.metadata)
|
||||
|
||||
if metadata.state != documentai.BatchProcessMetadata.State.SUCCEEDED:
|
||||
raise ValueError(f"Batch Process Failed: {metadata.state_message}")
|
||||
|
||||
storage_client = storage.Client()
|
||||
|
||||
# One process per Input Document
|
||||
for process in list(metadata.individual_process_statuses):
|
||||
# output_gcs_destination format: gs://BUCKET/PREFIX/OPERATION_NUMBER/INPUT_FILE_NUMBER/
|
||||
# The Cloud Storage API requires the bucket name and URI prefix separately
|
||||
matches = re.match(r"gs://(.*?)/(.*)", process.output_gcs_destination)
|
||||
if not matches:
|
||||
print(
|
||||
"Could not parse output GCS destination:",
|
||||
process.output_gcs_destination,
|
||||
)
|
||||
continue
|
||||
|
||||
output_bucket, output_prefix = matches.groups()
|
||||
|
||||
# Get List of Document Objects from the Output Bucket
|
||||
output_blobs = storage_client.list_blobs(output_bucket, prefix=output_prefix)
|
||||
|
||||
return list(output_blobs)
|
||||
|
||||
|
||||
def split_document(_doc):
|
||||
"""Splits a LangChain Document into smaller chunks."""
|
||||
# Use a recursive splitter for better semantic chunking
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=9216,
|
||||
chunk_overlap=200, # Add overlap for context
|
||||
)
|
||||
|
||||
new_docs = []
|
||||
for page in _doc:
|
||||
# Create smaller documents
|
||||
new_chunks = splitter.create_documents([page.page_content], [page.metadata])
|
||||
|
||||
# Reconstruct documents with the same metadata
|
||||
for i in range(len(new_chunks)):
|
||||
new_chunks[i].metadata["page_chunk"] = i
|
||||
new_chunks[i].metadata["chunk_size"] = len(new_chunks[i].page_content)
|
||||
new_docs.append(
|
||||
Document(
|
||||
page_content=new_chunks[i].page_content,
|
||||
metadata=new_chunks[i].metadata,
|
||||
)
|
||||
)
|
||||
|
||||
return new_docs
|
||||
|
||||
|
||||
# Triggered by a change in a storage bucket
|
||||
@functions_framework.cloud_event
|
||||
def process_pdf(cloud_event):
|
||||
"""Main function"""
|
||||
data = cloud_event.data
|
||||
|
||||
event_id = cloud_event["id"]
|
||||
event_type = cloud_event["type"]
|
||||
|
||||
bucket = data["bucket"]
|
||||
name = data["name"]
|
||||
metageneration = data["metageneration"]
|
||||
timeCreated = data["timeCreated"]
|
||||
updated = data["updated"]
|
||||
|
||||
print(f"Event ID: {event_id}")
|
||||
print(f"Event type: {event_type}")
|
||||
print(f"Bucket: {bucket}")
|
||||
print(f"File: {name}")
|
||||
print(f"Metageneration: {metageneration}")
|
||||
print(f"Created: {timeCreated}")
|
||||
print(f"Updated: {updated}")
|
||||
|
||||
file_name, extension = os.path.splitext(name)
|
||||
if extension != ".pdf":
|
||||
print("File is not a PDF. Please submit a PDF for processing instead.")
|
||||
return
|
||||
|
||||
# Project vars
|
||||
region = os.environ["REGION"]
|
||||
project_id = os.environ["PROJECT_ID"]
|
||||
|
||||
# Document AI Vars
|
||||
source_file = f"gs://{bucket}/{name}"
|
||||
gcs_output_uri = f"gs://{project_id}-doc-ai/doc-ai-output/" # Must end with a trailing slash `/`. Format: gs://bucket/directory/subdirectory/
|
||||
location = "us" # Format is "us" or "eu"
|
||||
processor_id = os.environ["PROCESSOR_ID"] # Create processor before running sample
|
||||
|
||||
blobs = batch_process_documents(
|
||||
project_id=project_id,
|
||||
location=location,
|
||||
processor_id=processor_id,
|
||||
gcs_output_uri=gcs_output_uri,
|
||||
gcs_input_uri=source_file, # Format: gs://bucket/directory/file.pdf
|
||||
input_mime_type="application/pdf",
|
||||
)
|
||||
|
||||
# Document AI may output multiple JSON files per source file
|
||||
lc_doc = []
|
||||
for blob in blobs:
|
||||
# Document AI should only output JSON files to GCS
|
||||
if blob.content_type != "application/json":
|
||||
print(
|
||||
f"Skipping non-supported file: {blob.name} - Mimetype: {blob.content_type}"
|
||||
)
|
||||
continue
|
||||
|
||||
# Download JSON File as bytes object and convert to Document Object
|
||||
print(f"Fetching {blob.name}")
|
||||
document = documentai.Document.from_json(
|
||||
blob.download_as_bytes(), ignore_unknown_fields=True
|
||||
)
|
||||
|
||||
# Create LangChain doc
|
||||
page = Document(
|
||||
page_content=document.text,
|
||||
metadata={
|
||||
"source": source_file,
|
||||
"page": document.shard_info.shard_index + 1,
|
||||
"ticker": Path(source_file).stem,
|
||||
"page_size": len(document.text),
|
||||
"doc_ai_shard_count": document.shard_info.shard_count,
|
||||
"doc_ai_shard_index": document.shard_info.shard_index,
|
||||
"doc_ai_chunk_size": blob._CHUNK_SIZE_MULTIPLE,
|
||||
"doc_ai_chunk_uri": blob.public_url,
|
||||
},
|
||||
)
|
||||
lc_doc.append(page)
|
||||
|
||||
# Split docs into smaller chunks (max 3072 tokens, 9216 characters)
|
||||
lc_doc_chunks = split_document(lc_doc)
|
||||
|
||||
# Setup embeddings
|
||||
embedding = VertexAIEmbeddings(model_name="text-embedding-005", project=project_id)
|
||||
|
||||
# AlloyDB Vars
|
||||
cluster = "alloydb-cluster"
|
||||
instance = "alloydb-instance"
|
||||
database = "ragdemos"
|
||||
table_name = "langchain_vector_store"
|
||||
user = "postgres"
|
||||
password = os.environ["ALLOYDB_PASSWORD"]
|
||||
initialize_vector_store = False
|
||||
ip_type = os.environ["IP_TYPE"]
|
||||
|
||||
# Create vector store
|
||||
engine = AlloyDBEngine.from_instance(
|
||||
project_id=project_id,
|
||||
region=region,
|
||||
cluster=cluster,
|
||||
instance=instance,
|
||||
database=database,
|
||||
user=user,
|
||||
password=password,
|
||||
ip_type=ip_type,
|
||||
)
|
||||
|
||||
if initialize_vector_store:
|
||||
engine.init_vectorstore_table(
|
||||
table_name=table_name,
|
||||
vector_size=768, # Vector size for VertexAI model(text-embedding-005)
|
||||
metadata_columns=[
|
||||
Column("source", "VARCHAR", nullable=True),
|
||||
Column("page", "INTEGER", nullable=True),
|
||||
Column("ticker", "VARCHAR", nullable=True),
|
||||
Column("page_size", "INTEGER", nullable=True),
|
||||
Column("doc_ai_shard_count", "INTEGER", nullable=True),
|
||||
Column("doc_ai_shard_index", "INTEGER", nullable=True),
|
||||
Column("doc_ai_chunk_size", "INTEGER", nullable=True),
|
||||
Column("doc_ai_chunk_uri", "VARCHAR", nullable=True),
|
||||
Column("page_chunk", "INTEGER", nullable=True),
|
||||
Column("chunk_size", "INTEGER", nullable=True),
|
||||
],
|
||||
overwrite_existing=True,
|
||||
)
|
||||
|
||||
store = AlloyDBVectorStore.create_sync(
|
||||
engine=engine,
|
||||
table_name=table_name,
|
||||
embedding_service=embedding,
|
||||
metadata_columns=[
|
||||
"source",
|
||||
"page",
|
||||
"ticker",
|
||||
"page_size",
|
||||
"doc_ai_shard_count",
|
||||
"doc_ai_shard_index",
|
||||
"doc_ai_chunk_size",
|
||||
"doc_ai_chunk_uri",
|
||||
"page_chunk",
|
||||
"chunk_size",
|
||||
],
|
||||
)
|
||||
|
||||
ids = [str(uuid.uuid4()) for i in range(len(lc_doc_chunks))]
|
||||
store.add_documents(lc_doc_chunks, ids)
|
||||
|
||||
print("Finished processing pdf")
|
||||
|
||||
# Send message to pubsub topic to kick off next step
|
||||
ticker = Path(source_file).stem
|
||||
publisher = pubsub_v1.PublisherClient()
|
||||
topic_name = f"projects/{project_id}/topics/{project_id}-doc-ready"
|
||||
future = publisher.publish(topic_name, bytes(f"{ticker}".encode()), spam="done")
|
||||
future.result()
|
||||
print("Sent message to pubsub")
|
||||
@@ -0,0 +1,11 @@
|
||||
functions-framework==3.*
|
||||
google-api-core==2.19.2
|
||||
google-cloud-documentai==2.32.0
|
||||
google-cloud-core==2.4.1
|
||||
google-cloud-pubsub==2.23.0
|
||||
google-cloud-storage==2.18.2
|
||||
langchain==0.3.15
|
||||
langchain-core==0.3.31
|
||||
langchain-google-alloydb-pg==0.9.3
|
||||
langchain-google-vertexai==2.0.15
|
||||
langchain-text-splitters==0.3.9
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Function to update the Vertex AI Search and Conversion index"""
|
||||
|
||||
import os
|
||||
|
||||
import functions_framework
|
||||
from google.api_core.client_options import ClientOptions
|
||||
from google.cloud import discoveryengine
|
||||
|
||||
|
||||
def import_documents_sample(
|
||||
project_id: str,
|
||||
location: str,
|
||||
data_store_id: str,
|
||||
gcs_uri: str | None = None,
|
||||
bigquery_dataset: str | None = None,
|
||||
bigquery_table: str | None = None,
|
||||
) -> str:
|
||||
"""Function to import documents"""
|
||||
# For more information, refer to:
|
||||
# https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store
|
||||
client_options = (
|
||||
ClientOptions(api_endpoint=f"{location}-discoveryengine.googleapis.com")
|
||||
if location != "global"
|
||||
else None
|
||||
)
|
||||
|
||||
# Create a client
|
||||
client = discoveryengine.DocumentServiceClient(client_options=client_options)
|
||||
|
||||
# The full resource name of the search engine branch.
|
||||
# e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}
|
||||
parent = client.branch_path(
|
||||
project=project_id,
|
||||
location=location,
|
||||
data_store=data_store_id,
|
||||
branch="default_branch",
|
||||
)
|
||||
|
||||
if gcs_uri:
|
||||
request = discoveryengine.ImportDocumentsRequest(
|
||||
parent=parent,
|
||||
gcs_source=discoveryengine.GcsSource(
|
||||
input_uris=[gcs_uri], data_schema="document"
|
||||
),
|
||||
# Options: `FULL`, `INCREMENTAL`
|
||||
reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
|
||||
)
|
||||
else:
|
||||
request = discoveryengine.ImportDocumentsRequest(
|
||||
parent=parent,
|
||||
bigquery_source=discoveryengine.BigQuerySource(
|
||||
project_id=project_id,
|
||||
dataset_id=bigquery_dataset,
|
||||
table_id=bigquery_table,
|
||||
data_schema="custom",
|
||||
),
|
||||
# Options: `FULL`, `INCREMENTAL`
|
||||
reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
|
||||
)
|
||||
|
||||
# Make the request
|
||||
operation = client.import_documents(request=request)
|
||||
|
||||
print(f"Waiting for operation to complete: {operation.operation.name}")
|
||||
response = operation.result()
|
||||
|
||||
# Once the operation is complete,
|
||||
# get information from operation metadata
|
||||
metadata = discoveryengine.ImportDocumentsMetadata(operation.metadata)
|
||||
|
||||
# Handle the response
|
||||
print(response)
|
||||
print(metadata)
|
||||
|
||||
return operation.operation.name
|
||||
|
||||
|
||||
@functions_framework.cloud_event
|
||||
def update_search_index(cloud_event):
|
||||
"""Main function"""
|
||||
# Set event vars
|
||||
data = cloud_event.data
|
||||
event_id = cloud_event["id"]
|
||||
event_type = cloud_event["type"]
|
||||
bucket = data["bucket"]
|
||||
name = data["name"]
|
||||
metageneration = data["metageneration"]
|
||||
timeCreated = data["timeCreated"]
|
||||
updated = data["updated"]
|
||||
|
||||
# Print event vars to log
|
||||
print(f"Event ID: {event_id}")
|
||||
print(f"Event type: {event_type}")
|
||||
print(f"Bucket: {bucket}")
|
||||
print(f"File: {name}")
|
||||
print(f"Metageneration: {metageneration}")
|
||||
print(f"Created: {timeCreated}")
|
||||
print(f"Updated: {updated}")
|
||||
|
||||
project_id = os.environ["PROJECT_ID"]
|
||||
location = "global"
|
||||
data_store_id = os.environ["DATASTORE_ID"]
|
||||
docs_metadata_bucket = os.environ["DOCS_METADATA_BUCKET"]
|
||||
gcs_uri = f"gs://{docs_metadata_bucket}/*.jsonl"
|
||||
|
||||
import_documents_sample(
|
||||
project_id=project_id,
|
||||
location=location,
|
||||
data_store_id=data_store_id,
|
||||
gcs_uri=gcs_uri,
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
functions-framework==3.*
|
||||
google-cloud-discoveryengine==0.12.2
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Function to write metadata for new pdf documents"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
import uuid
|
||||
|
||||
import functions_framework
|
||||
from google.cloud import storage
|
||||
|
||||
|
||||
@functions_framework.cloud_event
|
||||
def write_metadata(cloud_event):
|
||||
"""Main function"""
|
||||
# Set event vars
|
||||
data = cloud_event.data
|
||||
event_id = cloud_event["id"]
|
||||
event_type = cloud_event["type"]
|
||||
bucket = data["bucket"]
|
||||
name = data["name"]
|
||||
metageneration = data["metageneration"]
|
||||
timeCreated = data["timeCreated"]
|
||||
updated = data["updated"]
|
||||
|
||||
# Print event vars to log
|
||||
print(f"Event ID: {event_id}")
|
||||
print(f"Event type: {event_type}")
|
||||
print(f"Bucket: {bucket}")
|
||||
print(f"File: {name}")
|
||||
print(f"Metageneration: {metageneration}")
|
||||
print(f"Created: {timeCreated}")
|
||||
print(f"Updated: {updated}")
|
||||
|
||||
# Set local vars
|
||||
project_id = os.environ["PROJECT_ID"]
|
||||
metadata_bucket = f"{project_id}-docs-metadata"
|
||||
uid = str(uuid.uuid4())
|
||||
ticker = Path(name).stem
|
||||
target_file_name = f"{ticker}.jsonl"
|
||||
|
||||
# Build metadata jsonl
|
||||
metadata = (
|
||||
'{"id":"'
|
||||
+ uid
|
||||
+ '","structData":{"ticker":"'
|
||||
+ ticker
|
||||
+ '"},"content":{"mimeType":"application/pdf","uri":"gs://'
|
||||
+ bucket
|
||||
+ "/"
|
||||
+ name
|
||||
+ '"}}'
|
||||
)
|
||||
|
||||
# Write jsonl to metadata bucket
|
||||
storage_client = storage.Client()
|
||||
bucket = storage_client.bucket(metadata_bucket)
|
||||
blob = bucket.blob(target_file_name)
|
||||
|
||||
with blob.open("w") as f:
|
||||
f.write(metadata)
|
||||
|
||||
print(f"Metadata for {ticker} written to {metadata_bucket}")
|
||||
print(f"Metadata object: {metadata}")
|
||||
storage_client.close()
|
||||
@@ -0,0 +1,2 @@
|
||||
functions-framework==3.*
|
||||
google-cloud-storage==2.18.2
|
||||
Reference in New Issue
Block a user