chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 13:30:30 +08:00
commit 914fea506e
2793 changed files with 802106 additions and 0 deletions
@@ -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