Code Generation

pip install google-cloud-aiplatform langchain langchain-google-vertexai \
    langchain-community faiss-cpu nbformat

import vertexai
from langchain_google_vertexai import VertexAI, VertexAIEmbeddings

PROJECT_ID = "your-project-id"
LOCATION = "us-central1"

vertexai.init(project=PROJECT_ID, location=LOCATION)

# Initialize Gemini model
code_llm = VertexAI(
    model_name="gemini-2.0-flash",
    max_output_tokens=2048,
    temperature=0.1,
)

import requests
import nbformat
from langchain.schema.document import Document

def extract_python_code_from_ipynb(github_url):
    """Extract code cells from Jupyter notebooks."""
    raw_url = github_url.replace("github.com", "raw.githubusercontent.com").replace(
        "/blob/", "/"
    )
    
    response = requests.get(raw_url)
    response.raise_for_status()
    
    notebook = nbformat.reads(response.text, as_version=nbformat.NO_CONVERT)
    
    python_code = ""
    for cell in notebook.cells:
        if cell.cell_type == "code":
            python_code += cell.source + "\n"
    
    return python_code

# Crawl repository
def crawl_github_repo(repo: str, access_token: str):
    """Recursively list files in a GitHub repository."""
    api_url = f"https://api.github.com/repos/{repo}/contents"
    headers = {
        "Accept": "application/vnd.github.v3+json",
        "Authorization": f"Bearer {access_token}",
    }
    
    response = requests.get(api_url, headers=headers)
    response.raise_for_status()
    
    files = []
    for item in response.json():
        if item["type"] == "file" and item["name"].endswith(".ipynb"):
            files.append(item["html_url"])
        elif item["type"] == "dir":
            files.extend(crawl_github_repo(item["url"], access_token))
    
    return files

# Extract code from all notebooks
GITHUB_REPO = "GoogleCloudPlatform/generative-ai"
code_files_urls = crawl_github_repo(GITHUB_REPO, GITHUB_TOKEN)

code_documents = []
for i, url in enumerate(code_files_urls):
    if url.endswith(".ipynb"):
        content = extract_python_code_from_ipynb(url)
        doc = Document(
            page_content=content,
            metadata={"url": url, "file_index": i}
        )
        code_documents.append(doc)

from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS

# Chunk code intelligently
text_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.PYTHON,
    chunk_size=2000,
    chunk_overlap=200
)

texts = text_splitter.split_documents(code_documents)

# Create embeddings
embeddings = VertexAIEmbeddings(
    model_name="text-embedding-005",
    requests_per_minute=100,
)

# Build vector index
db = FAISS.from_documents(texts, embeddings)

# Initialize retriever
retriever = db.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5},
)

from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

rag_prompt = """
You are a proficient Python developer. Generate syntactically correct and concise code.

Rules:
1. Use the context to understand APIs and apply them correctly
2. Do not add license information to the output
3. Do not include Colab-specific code
4. Ensure all requirements in the question are met

Question:
{question}

Context:
{context}

Code:
"""

prompt_template = PromptTemplate(
    template=rag_prompt,
    input_variables=["context", "question"]
)

qa_chain = RetrievalQA.from_llm(
    llm=code_llm,
    prompt=prompt_template,
    retriever=retriever,
    return_source_documents=True,
)

user_question = """
Create a Python function that takes a prompt and generates a prediction 
using LangChain with the Vertex AI Gemini model.
"""

results = qa_chain.invoke(input={"query": user_question})
print(results["result"])

def predict_with_gemini(prompt: str) -> str:
    """Generate prediction using Vertex AI Gemini via LangChain.
    
    Args:
        prompt: Input text prompt
        
    Returns:
        Model's text response
    """
    from langchain_google_vertexai import VertexAI
    import vertexai
    
    # Initialize Vertex AI
    vertexai.init(project="your-project", location="us-central1")
    
    # Create model instance
    llm = VertexAI(
        model_name="gemini-2.0-flash",
        max_output_tokens=2048,
        temperature=0.1,
    )
    
    # Generate response
    response = llm.invoke(input=prompt)
    
    return response

# Usage
result = predict_with_gemini("Explain quantum computing")
print(result)

user_question = """
Create a Python function that takes text input and returns embeddings 
using LangChain with Vertex AI text-embedding-005 model.
"""

results = qa_chain.invoke(input={"query": user_question})
print(results["result"])

def get_embeddings(text: str | list[str]) -> list[list[float]]:
    """Generate embeddings for text using Vertex AI.
    
    Args:
        text: Single string or list of strings to embed
        
    Returns:
        List of embedding vectors
    """
    from langchain_google_vertexai import VertexAIEmbeddings
    import vertexai
    
    vertexai.init(project="your-project", location="us-central1")
    
    embeddings = VertexAIEmbeddings(
        model_name="text-embedding-005",
    )
    
    if isinstance(text, str):
        text = [text]
    
    vectors = embeddings.embed_documents(text)
    
    return vectors

# Usage
embeddings = get_embeddings(["Hello world", "AI is amazing"])
print(f"Generated {len(embeddings)} embeddings")

# Without context
zero_shot_response = code_llm.invoke(
    input=user_question,
    max_output_tokens=2048,
    temperature=0.1
)

print("Zero-shot output:")
print(zero_shot_response)

# With codebase context
rag_response = qa_chain.invoke(input={"query": user_question})

print("RAG-enhanced output:")
print(rag_response["result"])

print("\nSource documents:")
for doc in rag_response["source_documents"]:
    print(f"- {doc.metadata['url']}")

from google import genai
from google.genai.types import GenerateContentConfig

client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

legacy_code = """
def process_data(data):
    result = []
    for i in range(len(data)):
        if data[i] > 0:
            result.append(data[i] * 2)
    return result
"""

refactoring_prompt = f"""
Refactor this code to use modern Python best practices:
- Use list comprehensions
- Add type hints
- Add docstrings
- Use descriptive variable names

Code:
{legacy_code}
"""

response = client.models.generate_content(
    model="gemini-2.0-flash",
    contents=refactoring_prompt,
    config=GenerateContentConfig(temperature=0.1),
)

print(response.text)

def process_positive_numbers(data: list[int | float]) -> list[int | float]:
    """Double all positive numbers in the input list.
    
    Args:
        data: List of numeric values
        
    Returns:
        List containing doubled positive values
    """
    return [value * 2 for value in data if value > 0]

undocumented_code = """
def calculate_metrics(transactions, start_date, end_date):
    filtered = [t for t in transactions if start_date <= t['date'] <= end_date]
    total = sum(t['amount'] for t in filtered)
    avg = total / len(filtered) if filtered else 0
    return {'total': total, 'average': avg, 'count': len(filtered)}
"""

doc_prompt = f"""
Add comprehensive docstrings to this function following Google style:

{undocumented_code}
"""

response = client.models.generate_content(
    model="gemini-2.0-flash",
    contents=doc_prompt,
)

print(response.text)

suspicious_code = """
def execute_query(user_input):
    query = f"SELECT * FROM users WHERE username = '{user_input}'"
    return database.execute(query)
"""

security_prompt = f"""
Analyze this code for security vulnerabilities:

{suspicious_code}

Provide:
1. List of vulnerabilities
2. Severity rating
3. Remediation code
"""

response = client.models.generate_content(
    model="gemini-2.0-flash",
    contents=security_prompt,
)

print(response.text)