Files
schwartz1375-genai-essentials/llm_tutorial.ipynb
T
schwartz1375 6c257ebab0 Change embedding model to embeddinggemma and fix ChromaDB initialization
- Switch from nomic-embed-text to embeddinggemma for embeddings
- Add directory verification cells to ensure required directories exist
- Add error handling for corrupted ChromaDB databases with auto-recovery
- Create persist directories before ChromaDB initialization
2025-12-06 16:59:15 -05:00

99 KiB

Environment Setup

This notebook demonstrates LangChain concepts using local models via Ollama. Before continuing, make sure the required models and Python packages are installed.


1. Install and Configure Ollama

If you haven't already, install Ollama and ensure it's running locally.

Then pull the required models:

ollama pull gemma3
ollama pull embeddinggemma

These models will be used for generation (gemma3) and embeddings (embeddinggemma).


2. Install Required Python Packages

In [1]:
'''
!pip install --upgrade pip --quiet
!pip install -U \
  langchain langchain-ollama langchain-community langchain-core \
  langchain-experimental chromadb pandas pymupdf ipython \
  duckduckgo-search --quiet
'''
Out [1]:
'\n!pip install --upgrade pip --quiet\n!pip install -U   langchain langchain-ollama langchain-community langchain-core   langchain-experimental chromadb pandas pymupdf ipython   duckduckgo-search --quiet\n'

3. Start Ollama (if not already running)

Make sure Ollama is running in the background. You can check by running:

ollama list

This should list both gemma3:latest and embeddinggemma models.

Local vs. Cloud Models

Ollama (Local Execution)

Ollama allows you to run open-source models (e.g., gemma, llama, embeddinggemma) locally. This provides control, privacy, and fast iteration without API latency or cost.

Amazon Bedrock (Cloud Execution)

Amazon Bedrock provides access to multiple foundation models through a single API. It is fully managed, scalable, and integrates into AWS’s security and governance ecosystem.

Basic Concepts of Generative AI

This section introduces core concepts for working with Large Language Models (LLMs), including how to prompt them effectively and build applications using frameworks like LangChain—whether deploying models locally via Ollama or through cloud services like AWS Bedrock.


What Are LLMs?

Large Language Models (LLMs) are machine learning models trained on massive text corpora. They predict the next token in a sequence based on context, making them powerful pattern-recognition tools for generating human-like text. However, they do not understand meaning or possess reasoning capabilities—their output is probabilistic, not factual by default.


Prompting and Prompt Engineering

Prompts are structured inputs that tell an LLM what to do. A well-crafted prompt can significantly influence the quality, accuracy, and usefulness of the model's response.

Prompts generally consist of three key components:


Instruction

The instruction specifies the task or behavior expected from the model. It can be a simple question, a directive, or even a persona setup.

Examples:

  • “Generate a list of known malicious domains associated with the Qakbot malware family.”
  • “Write a poem about the Qakbot malware family.”
  • “Produce a sonnet about my love of Franklin’s BBQ.”

Context

Context includes relevant information provided to the model to help it complete the instruction. This is particularly important for specialized or domain-specific use cases.

Examples of context:

  • Extracted text from a PDF or website
  • Lists of known facts or structured datasets
  • Internal business knowledge

Context is critical when the model lacks the necessary background knowledge. It also enables in-context learning and is foundational to RAG (Retrieval-Augmented Generation) systems.

Reference:


Output Format

Defining the output format ensures the model returns a response in a structure that's compatible with downstream processing or presentation.

Examples:

  • JSON/JSONL for programmatic consumption
  • Markdown for documentation or emails
  • Lists, CSV, or tabular formats for parsing

Providing formatting examples in your prompt often leads to more consistent and usable output.


Best Practices for Prompt Engineering

  • Be explicit and unambiguous in your instructions.
  • Focus on what you want, not just what to avoid.
  • Supply sufficient context for the model to reason accurately.
  • Specify format, tone, language, or response length as needed.
  • Iterate and refine—prompt design is an interactive process.
  • Be cost-conscious—longer prompts and responses mean higher token usage.
  • Leverage in-context examples to teach the model new behavior.

Prompting strategies:

  • Zero-shot: No examples provided.
  • One-shot: One example included.
  • Few-shot: Multiple examples included.

System Prompts

Some platforms (e.g., OpenAI, Anthropic) support system prompts—instructions that shape the model’s persona, tone, or behavioral constraints.

Example system message:

“You are a cybersecurity analyst generating threat intelligence summaries.”

System prompts are not typically hidden; providers often document their usage:


Fine-Tuning vs. In-Context Learning

  • Fine-tuning modifies the model’s internal weights using domain-specific data. It’s powerful but resource-intensive and harder to maintain.
  • In-context learning gives examples inline in the prompt. It’s flexible, efficient, and sufficient for many real-world use cases.

Learn more: Prompt Engineering Guide

Core LangChain Concepts

LangChain is a modular framework for building generative AI applications powered by Large Language Models (LLMs). It simplifies the development of multi-step pipelines, retrieval-augmented workflows, and agent-based reasoning systems.


LangChain Capabilities

LangChain provides key building blocks for LLM applications:

  • Prompt templates for reusability and structure
  • Memory for multi-turn interactions
  • Tools to access external data sources
  • Chains to compose tasks in sequence or branching logic
  • Agents for dynamic tool use and reasoning
  • Document loaders, chunkers, and retrievers for knowledge ingestion

Prompt Templates

Prompt templates allow you to define reusable prompts with variables. This is essential for consistent prompt formatting in production use cases.

from langchain.prompts import PromptTemplate

template = "Generate a list of domains associated with {malware_family}"
prompt = PromptTemplate.from_template(template)
prompt.format(malware_family="Qakbot")

Memory

LLMs are inherently stateless, meaning they do not retain context between interactions. LangChain provides memory modules such as ConversationBufferMemory that persist prior messages and automatically re-inject them into prompts.

from langchain.memory import ConversationBufferMemory

memory = ConversationBufferMemory()

When used in a chain:

chain = LLMChain(llm=llm, prompt=prompt_template, memory=memory)

Memory is critical for applications like chatbots, copilots, or multi-turn reasoning agents.


Tools

Tools allow LLMs to interact with external systems—APIs, web search engines, internal databases, or third-party services. LangChain supports tool integration via descriptors that let agents reason about what tool to use and when.

Examples of tools:

  • Google/Bing search
  • DomainTools, VirusTotal, Wikipedia
  • Custom internal APIs (e.g., threat intelligence, ticketing systems)

More: LangChain Tools Documentation


Chains

Chains are workflows composed of modular steps. Each chain links one or more components (e.g., prompt → LLM → post-processor).

Common chain types include:

  • LLMChain: single prompt-response (Note: LLMChain is deprecated in favor of LCEL - LangChain Expression Language)
  • SimpleSequentialChain: sequential chaining of steps
  • MultiInputChain: accepts multiple inputs
  • RouterChain: dynamically selects downstream chains based on input

Example with memory:

chain = LLMChain(llm=llm, prompt=prompt_template, memory=ConversationBufferMemory())

Chains provide structure and allow you to build reliable pipelines with conditional logic and state.


Agents

Agents enable dynamic, reasoning-driven workflows. Unlike static chains, agents interpret user input, determine which tools or sub-chains to invoke, gather additional data, and synthesize final responses.

They excel at handling ambiguous prompts, complex logic, or open-ended tasks.

More: LangChain Agents Documentation

RAG and Embeddings

Vector Stores

Vector databases (e.g., Chroma, FAISS, Pinecone) store text embeddings, which allow similarity searches against known documents or facts.

RAG (Retrieval-Augmented Generation)

RAG improves response quality by retrieving relevant content from an external knowledge source before prompting the model. This enables high-accuracy answers even from models that weren’t explicitly trained on the content.

Chunking

Long documents are split into smaller chunks to improve retrieval accuracy and avoid exceeding token limits. LangChain includes text splitters like RecursiveCharacterTextSplitter that optimize chunk size and overlap.

Embeddings

Embeddings are numeric vectors that represent semantic meaning. They are used for similarity search in vector stores.

  • Local: e.g., embeddinggemma via Ollama
  • Cloud: e.g., Titan via AWS Bedrock, HuggingFace SentenceTransformers

LLM Limitations

  • Hallucination: Models may fabricate plausible but incorrect information.
  • Context limitations: Only a fixed number of tokens can be processed per prompt.
  • Data leakage risk: LLMs trained on sensitive data may inadvertently reveal it.
  • Stateless: Models forget previous interactions unless provided memory mechanisms.

Understanding Tokenization

LLMs process input/output as tokens, not words. Token count affects:

  • Cost (in usage-based billing models)
  • Whether your full prompt fits in the model’s context window
  • Output truncation risk

Useful tools:

In [2]:
from IPython.display import Markdown
from langchain_ollama import ChatOllama 

# Note: Ensure Ollama is running locally and the gemma3 model is installed
# Run: ollama pull gemma3
# If Ollama is not running, you'll get a connection error

llm = ChatOllama(
    model="gemma3",
    temperature=0.8,
    # base_url="http://remote-ip:port", # Ollama default port: 11434
    # timeout=5,
    num_predict=256,  # Use num_predict instead of max_tokens for Ollama
    # other params ...
)

messages = [
    ("system", "You are a helpful assistant."),
    ("human", "why is the sky blue?"),
]

Markdown(llm.invoke(messages).content)
Out [2]:
<IPython.core.display.Markdown object>

Using Prompt Templates and LangChain Expression Syntax

This example introduces the use of a ChatPromptTemplate combined with an Ollama-backed LLM using LangChain's expression syntax. This pattern allows for reusable, parameterized prompts and demonstrates how to integrate structured prompting into your pipeline. The | operator chains the prompt and model together, creating a modular and declarative workflow.

We’ll also simulate a multi-role setup (e.g., system + human) within the prompt string.

In [3]:
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama.llms import OllamaLLM

template = """Question: {question}

Answer: Let's think step by step."""

prompt = ChatPromptTemplate.from_template(template)

model = OllamaLLM(model="gemma3")

chain = prompt | model

Markdown(chain.invoke({"question": "System: You are cybersecurity expert, and AI assisant to AWS.\n\nHuman: What is the LummaC2 malware?"}))
Out [3]:
<IPython.core.display.Markdown object>

Memory

It's nice to be able to ask a question and get an answer, but that seems pretty transactional and impersonal. I'd like the LLM I'm talking to to show me they are listening to me and paying attention to what I am saying. I am, after all, human and I need to be loved or at least feel like it.

But LLMs are cold and heartless (technically, they are stateless but we are talking about my feelings here). This means that they don’t ‘remember’ interactions from prompt to prompt. You can fine tune them to persist data but without updating the weights through expensive training but, they don’t ‘remember’ anything more after the model is set at a checkpoint. This means if I just asked what the Qakbot malware was and then follow up with a question like "What industry was primarily targeted by this malware?", the model will not be able to answer within the context of Qakbot...since show beats tell, let's see how that works firsthand.

Now lets ask the followup question...

In [4]:
Markdown(llm.invoke("System: You are cybersecurity expert, and AI assisant.\n\nHuman: What industry was primarily targeted by this malware?").content)
Out [4]:
<IPython.core.display.Markdown object>

LangChain Memory Components: Enabling Stateful Conversations

The memory component of LangChain allows the LLM to become 'stateful'. This quality is quite useful when developing applications driven by LLMs. For instance, a conversational system or a chatbot is required to recall past interactions to maintain a conversation. Without memory, the system would not be able to handle follow-up messages or recollect key pieces of information mentioned earlier in the conversation.

In this section, we will explore the memory modules provided by LangChain. LangChain offers several types of memory modules depending on the task and the properties of the LLM. We will incorporate memory into chains and examine changes in performance due to it.

In [5]:
from langchain_community.llms import Ollama
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
import pandas as pd
import warnings

warnings.filterwarnings("ignore")

# Initialize Ollama with Gemma model
llm = Ollama(model="gemma3")
In [6]:
# Set up conversation memory
cbm_memory = ConversationBufferMemory()

# Adding memory to a conversational chain
chain_with_buffer_memory = ConversationChain(llm=llm, memory=cbm_memory)

# Helper function to prompt and print memory state
def prompt_and_print_memory(prompts, chain_with_memory):
    # Store the responses
    responses = {"input": [], "history": [], "response": []}

    # Repeatedly prompting the chain and observing the memory
    for prompt in prompts:
        response = chain_with_memory.invoke({"input": prompt})
        responses["input"].append(prompt)
        responses["history"].append(cbm_memory.load_memory_variables({}))
        responses["response"].append(response["response"])

    # Display responses in a dataframe
    df = pd.DataFrame.from_dict(responses)
    with pd.option_context("display.max_colwidth", None):
        display(df)

# Sequence of prompts for demonstration
cbm_prompts = [
    "What is the LummaC2 malware?",
    "What industry was primarily targeted by this malware?",
]

# Invoke helper function
prompt_and_print_memory(cbm_prompts, chain_with_buffer_memory)
input history response
0 What is the LummaC2 malware? {'history': 'Human: What is the LummaC2 malware? AI: Okay, let’s talk about LummaC2! It’s a really interesting and quite sophisticated piece of malware that’s been tracked by analysts at Sophos. Essentially, it's a modular, multi-stage post-exploitation framework primarily targeting Windows systems. Here’s a breakdown of what makes it notable: * **Developed by the Chinese APT group, UNC3923:** Sophos has linked it very strongly to this group, which is known for its espionage activities, often targeting the energy, defense, and technology sectors. They’ve been tracking this group for a *long* time – back to 2018 – and LummaC2 is a key component of their toolkit. * **Modular Design:** This is a *big* deal. LummaC2 isn't just one piece of malware; it’s composed of several modules that can be used individually or combined. This makes it highly adaptable and difficult to detect. The main modules include: * **Beacon:** This is the core module. It’s a C2 (Command and Control) beacon that allows the attackers to communicate with the infected host. It uses TLS encryption for secure communication – they’re not sending data in the clear. The beacon itself utilizes multiple protocols – HTTP, HTTPS, and even DNS – to evade detection. * **Credential Dumping:** This module attempts to steal credentials (usernames, passwords, hashes) from the infected system. They've been observed using tools like Mimikatz – a very common tool for credential theft. * **Process Injection:** They use this to inject malicious code into legitimate running processes, further expanding their control over the system. * **Privilege Escalation:** LummaC2 attempts to escalate privileges on the infected system, allowing them to gain higher levels of control. * **Lateral Movement:** This is crucial. Once they have control of one machine, they use it to move laterally across the network, compromising other systems. They’ve been observed using tools like PsExec for this. * **TLS Encryption:** As I mentioned before, the beacon uses TLS (Transport Layer Security) encryption. This makes it harder for security analysts to monitor the communication and understand what the malware is doing in real-time. * **Observed Targets:** They’ve been actively targeting companies in the defense, energy, and technology sectors – specifically those with a global presence. * **Sophos Tracking:** Sophos has been diligently tracking LummaC2's activity since 2021, publishing detailed reports on its behavior and associated IOCs (Indicators of Compromise) – like IP addresses, domains, and file hashes – that security teams can use to detect and block the malware. They even have a dedicated threat feed for LummaC2. You can find a lot of detailed information on their website here: [https://www.sophos.com/threat-center/lummac2.html](https://www.sophos.com/threat-center/lummac2.html) Do you want me to delve deeper into a specific aspect of LummaC2, such as its modules, the IOCs, or how Sophos tracks it? Or perhaps you’d like to discuss how to defend against attacks from this group?'} Okay, let’s talk about LummaC2! It’s a really interesting and quite sophisticated piece of malware that’s been tracked by analysts at Sophos. Essentially, it's a modular, multi-stage post-exploitation framework primarily targeting Windows systems.\n\nHere’s a breakdown of what makes it notable:\n\n* **Developed by the Chinese APT group, UNC3923:** Sophos has linked it very strongly to this group, which is known for its espionage activities, often targeting the energy, defense, and technology sectors. They’ve been tracking this group for a *long* time – back to 2018 – and LummaC2 is a key component of their toolkit.\n\n* **Modular Design:** This is a *big* deal. LummaC2 isn't just one piece of malware; it’s composed of several modules that can be used individually or combined. This makes it highly adaptable and difficult to detect. The main modules include:\n * **Beacon:** This is the core module. It’s a C2 (Command and Control) beacon that allows the attackers to communicate with the infected host. It uses TLS encryption for secure communication – they’re not sending data in the clear. The beacon itself utilizes multiple protocols – HTTP, HTTPS, and even DNS – to evade detection.\n * **Credential Dumping:** This module attempts to steal credentials (usernames, passwords, hashes) from the infected system. They've been observed using tools like Mimikatz – a very common tool for credential theft.\n * **Process Injection:** They use this to inject malicious code into legitimate running processes, further expanding their control over the system.\n * **Privilege Escalation:** LummaC2 attempts to escalate privileges on the infected system, allowing them to gain higher levels of control.\n * **Lateral Movement:** This is crucial. Once they have control of one machine, they use it to move laterally across the network, compromising other systems. They’ve been observed using tools like PsExec for this.\n\n* **TLS Encryption:** As I mentioned before, the beacon uses TLS (Transport Layer Security) encryption. This makes it harder for security analysts to monitor the communication and understand what the malware is doing in real-time.\n\n* **Observed Targets:** They’ve been actively targeting companies in the defense, energy, and technology sectors – specifically those with a global presence.\n\n* **Sophos Tracking:** Sophos has been diligently tracking LummaC2's activity since 2021, publishing detailed reports on its behavior and associated IOCs (Indicators of Compromise) – like IP addresses, domains, and file hashes – that security teams can use to detect and block the malware. They even have a dedicated threat feed for LummaC2. You can find a lot of detailed information on their website here: [https://www.sophos.com/threat-center/lummac2.html](https://www.sophos.com/threat-center/lummac2.html)\n\nDo you want me to delve deeper into a specific aspect of LummaC2, such as its modules, the IOCs, or how Sophos tracks it? Or perhaps you’d like to discuss how to defend against attacks from this group?
1 What industry was primarily targeted by this malware? {'history': 'Human: What is the LummaC2 malware? AI: Okay, let’s talk about LummaC2! It’s a really interesting and quite sophisticated piece of malware that’s been tracked by analysts at Sophos. Essentially, it's a modular, multi-stage post-exploitation framework primarily targeting Windows systems. Here’s a breakdown of what makes it notable: * **Developed by the Chinese APT group, UNC3923:** Sophos has linked it very strongly to this group, which is known for its espionage activities, often targeting the energy, defense, and technology sectors. They’ve been tracking this group for a *long* time – back to 2018 – and LummaC2 is a key component of their toolkit. * **Modular Design:** This is a *big* deal. LummaC2 isn't just one piece of malware; it’s composed of several modules that can be used individually or combined. This makes it highly adaptable and difficult to detect. The main modules include: * **Beacon:** This is the core module. It’s a C2 (Command and Control) beacon that allows the attackers to communicate with the infected host. It uses TLS encryption for secure communication – they’re not sending data in the clear. The beacon itself utilizes multiple protocols – HTTP, HTTPS, and even DNS – to evade detection. * **Credential Dumping:** This module attempts to steal credentials (usernames, passwords, hashes) from the infected system. They've been observed using tools like Mimikatz – a very common tool for credential theft. * **Process Injection:** They use this to inject malicious code into legitimate running processes, further expanding their control over the system. * **Privilege Escalation:** LummaC2 attempts to escalate privileges on the infected system, allowing them to gain higher levels of control. * **Lateral Movement:** This is crucial. Once they have control of one machine, they use it to move laterally across the network, compromising other systems. They’ve been observed using tools like PsExec for this. * **TLS Encryption:** As I mentioned before, the beacon uses TLS (Transport Layer Security) encryption. This makes it harder for security analysts to monitor the communication and understand what the malware is doing in real-time. * **Observed Targets:** They’ve been actively targeting companies in the defense, energy, and technology sectors – specifically those with a global presence. * **Sophos Tracking:** Sophos has been diligently tracking LummaC2's activity since 2021, publishing detailed reports on its behavior and associated IOCs (Indicators of Compromise) – like IP addresses, domains, and file hashes – that security teams can use to detect and block the malware. They even have a dedicated threat feed for LummaC2. You can find a lot of detailed information on their website here: [https://www.sophos.com/threat-center/lummac2.html](https://www.sophos.com/threat-center/lummac2.html) Do you want me to delve deeper into a specific aspect of LummaC2, such as its modules, the IOCs, or how Sophos tracks it? Or perhaps you’d like to discuss how to defend against attacks from this group? Human: What industry was primarily targeted by this malware? AI: Okay, great question! While LummaC2 has been observed targeting a *range* of industries, the primary targets have consistently been the defense, energy, and technology sectors. As I mentioned before, Sophos has specifically noted a strong focus on companies with a global presence within those industries. They’ve seen a significant number of campaigns directed at organizations involved in defense technology – think things like missile defense systems or cybersecurity products – and companies involved in energy production and distribution. It's not that other industries weren't *ever* targeted, but those three were the most frequent and impactful. Do you want me to elaborate on *why* these sectors were particularly attractive to UNC3923, or perhaps discuss the specific types of data they were seeking within those industries? Or would you like me to go back and explain a particular module in more detail, like the lateral movement capabilities?'} Okay, great question! While LummaC2 has been observed targeting a *range* of industries, the primary targets have consistently been the defense, energy, and technology sectors. As I mentioned before, Sophos has specifically noted a strong focus on companies with a global presence within those industries. They’ve seen a significant number of campaigns directed at organizations involved in defense technology – think things like missile defense systems or cybersecurity products – and companies involved in energy production and distribution. It's not that other industries weren't *ever* targeted, but those three were the most frequent and impactful.\n\nDo you want me to elaborate on *why* these sectors were particularly attractive to UNC3923, or perhaps discuss the specific types of data they were seeking within those industries? Or would you like me to go back and explain a particular module in more detail, like the lateral movement capabilities?

Tools

We'll use a tool supplied with LangChain and a custom tool we create

DuckDuckGo is an internet privacy company most popularly known for their private search engine. The company emphasizes privacy and anonimity as one of the key principles behind all their products.

Let's create a DuckDuckGo tool that is capable of retrieving results from a web search.

In [7]:
from langchain.tools import Tool
from langchain_community.tools import DuckDuckGoSearchRun

# Define the API wrapper for DuckDuckGo search
duckduckgo_search = DuckDuckGoSearchRun()

# Define the DuckDuckGo tool using a description and the function to retrieve results from DuckDuckGo search
duckduckgo_tool = Tool(
    name="DuckDuckGoSearch",
    func=duckduckgo_search.run,
    description="useful for when you need to answer questions about current weather and other current events",
)

# Test the DuckDuckGo tool
Markdown(duckduckgo_tool("What is a honeybee?"))
Out [7]:
<IPython.core.display.Markdown object>

Custom Tools

You can define custom tools using the tool decorator.

It is critical when writing custom tools that we use well written doc strings, that we use python's type decoration syntax and that we use the @tool decorator supplied by LangChain (loaded above). The docstring and the type labels are used by the agent to reason about what a tool is useful for and how to interface with it. If your agnet isn't picking the correct tools or is calling them incorrectly, these are common culprits.

Keep in mind that a custom tool can do anything you can define in code. They can even call APIs of much larger systems. Effectively, these are boundless.

Here is a simple custom tool that returns the current date.

In [8]:
from langchain.tools import tool
from datetime import date

@tool
def curr_date(text: str) -> str:
    """Returns todays date, use this for any \
    questions related to knowing todays date. \
    The input should always be an empty string, \
    and this function will always return todays \
    date - any date math should occur \
    outside this function."""
    return str(date.today())
In [9]:

# Define the date tool
date_tool = Tool(
    name="DateTool", func=curr_date, description="Useful to retrieve the current date"
)

# Test date tool
print(date_tool(""))
2025-12-06

Agents

Finally...we get to the good stuff...Agents are really cool. Generally speaking, they have a model and a set of tools at their disposal. They can take a prompt, break it down into the steps necessary to produce a meaningful response and, using the descriptions of the tools at their disposal, decide which ones are likely to be able to accomplish each tasks in the best way. It then uses it's conclusions to execute the needed tools, pull in the right context and respond to your prompt.

They are simple to build and powerful to use.

Let's see if we can solve the dot counting problem now...

In [10]:
from langchain_experimental.tools import PythonREPLTool
from langchain.agents import AgentType
from langchain.agents import initialize_agent
from IPython.display import Markdown

llm = ChatOllama(
    model="gemma3",
    temperature=0.1,
    num_predict=256,  # Use num_predict instead of max_tokens for Ollama
)

# Define the Python REPL tool
python_repl = PythonREPLTool()

# Initialize the agent using the local model and the Python REPL tool
python_agent = initialize_agent(
    [python_repl],
    llm,  # Using our previously initialized Ollama Gemma model here
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True,
)

# Test the Python Agent with a simple input
result = python_agent.invoke(
    {"input": "How many dots are in the string ......................................?"}
)

# Display the result using Markdown
Markdown(result['output'])
Out [10]:

> Entering new AgentExecutor chain...
Python REPL can execute arbitrary code. Use with caution.
Thought: The string contains a series of dots. I need to count the number of dots.
Action: Python_REPL
Action Input: `string = "....................................."; print(len(string))`
Observation: 37

Thought:Thought: I now know the final answer
Final Answer: 37

> Finished chain.
<IPython.core.display.Markdown object>

Pulling it all together and let the model reason about something to get the answer we need.

Agent with Automatic Tool Selection

The real power of agents is when they are allowed to reason about which tools can be used to solve different tasks or sub-tasks to respond to a prompt. This is what makes them seem intelligent...but remember they are not truly intelligent! They're following sophisticated pattern recognition.

The agent reasoning flow works like this:

User Input → The user submits a query or request LLM Reasoning → The LLM analyzes the request and determines approach Tool Selection → Based on tool descriptions, the LLM selects appropriate tool(s) Tool Execution → The selected tools are executed with parameters from the LLM Result Synthesis → The LLM combines tool outputs into a coherent response Hopefully, the model decides correctly and you get great results. Sometimes you have to prompt the model to use the right tools.

For example, if a model is trying to solve a math problem without using your math-specific tool, you can just add "You are bad at math so always use the wolfram tool to solve math problems" in the prompt. Simple guidance like this is usually enough to get the model working the way you intend.

This reasoning ability to select appropriate tools creates a flexible system that can handle diverse queries by combining specialized tools with general language understanding. The agent can break down complex problems, identify which tools would be most helpful for each component, and then integrate the results into a coherent response.

In [11]:
from langchain.agents import initialize_agent

# Define a list of available tools for the agent
tools = [
    duckduckgo_tool,
    date_tool,
]

# Initialize the agent with access to all the tooks in the list
agent_executor = initialize_agent(
    tools, 
    llm, 
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, 
    verbose=True, 
    handle_parsing_errors=True,
)
In [12]:
agent_executor.invoke(
    {
        "input": """
        System: You are not good at determining dates and times.  When you are asked about weather reports, you should first make sure you know what date is being asked about.
        If you are asked about weather and a day was not specified, you should use the date_tool to get today's date and then use the duckduckgo_tool to find out the weather for the date.
        When asking about weather, you should always include the date in your search for weather so you get the weather report for the correct date.
        
        Human: I am in Washington DC.  Will I need an umbrella today?
        """
    }
)
Out [12]:

> Entering new AgentExecutor chain...
I need to find out the weather in Washington DC today to determine if I need an umbrella. Since the date was not specified, I should use the DateTool to get today's date and then use the DuckDuckGoSearch tool to find the weather for that date.
Action: DateTool
Action Input: ''
Observation: 2025-12-06
Thought:Okay, I now have today's date as 2025-12-06. I will use the DuckDuckGoSearch tool to find the weather in Washington DC for this date.
Action: DuckDuckGoSearch
Action Input: 'weather Washington DC 2025-12-06'
Observation: In Washington , a combination of overcast , rain-soaked and dry and sunny weather is expected for the next ten days. On Friday, in Washington , sunrise will be at 7:10 am and sunset at 4:47 pm (EST). ... Updated: Dec 4, 2025 @ 22:17:39 UTC Published by: Weather U.S. Washington , DC 39°F Cloudy- Light rain is expected. ... Delightful, is there generally a big difference in weather between S Scotland and the ... The Weather Channel uses data, cookies and other similar technologies on this browser to optimise our website, and to provide you with weather ... For example, the weather in Washington DC in April 2025 . ... Do you see something missing or inaccurate about the weather in Washington DC ? Feel ...
Thought:I now know the weather in Washington DC on 2025-12-06 is Cloudy- Light rain is expected.
Final Answer: Yes, you should probably take an umbrella with you today.

> Finished chain.
{'input': "\n        System: You are not good at determining dates and times.  When you are asked about weather reports, you should first make sure you know what date is being asked about.\n        If you are asked about weather and a day was not specified, you should use the date_tool to get today's date and then use the duckduckgo_tool to find out the weather for the date.\n        When asking about weather, you should always include the date in your search for weather so you get the weather report for the correct date.\n\n        Human: I am in Washington DC.  Will I need an umbrella today?\n        ",
 'output': 'Yes, you should probably take an umbrella with you today.'}

RAG Demonstration with ChromaDB

This cell demonstrates building a Retrieval-Augmented Generation (RAG) system using ChromaDB as the vector store. The code loads PDF documents, splits them into manageable chunks, creates vector embeddings using AWS Bedrock or HuggingFace models, and stores them in a persistent ChromaDB database. It showcases two key RAG functionalities: (1) direct similarity search to retrieve relevant document chunks based on semantic meaning, and (2) question answering that leverages the retrieved context to generate accurate responses using an LLM. This implementation includes fallback options, proper error handling, and performance testing to ensure reliable operation even if the primary embedding service encounters issues. The system persists the vector database to disk, allowing for reuse without regenerating embeddings in future sessions.

In [13]:
import os
import pandas as pd
from IPython.display import Markdown, display

# LangChain imports
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OllamaEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_community.llms import Ollama
from chromadb.config import Settings
In [14]:
from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from langchain_community.document_loaders import PyMuPDFLoader

import os
from IPython.display import Markdown, display
from langchain_core.documents import Document

# Initialize Ollama LLM (Gemma model)
# Make sure these are pulled once:
#   ollama pull gemma3
#   ollama pull nomic-embed-text
llm = Ollama(
    model="gemma3",
    base_url="http://localhost:11434",     # explicit, just to be clear
)

# 🔧 Embeddings model: add num_ctx and base_url
embedding_model = OllamaEmbeddings(
    model="embeddinggemma",
    base_url="http://localhost:11434",
    num_ctx=2048,     # <= keep context modest to avoid crashes
    # you could even try 1024 if needed
)

# Create directory for PDFs
pdf_dir = "./data/pdfs"
os.makedirs(pdf_dir, exist_ok=True)
print(f"PDF directory is at: {os.path.abspath(pdf_dir)}")

# 🔧 Text splitter: smaller chunks
recur_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,        # was 1000
    chunk_overlap=100,     # was 200
    separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""],
    length_function=len
)

# Load and process PDFs using PyMuPDFLoader for each file in the directory
data = []
for filename in os.listdir(pdf_dir):
    if filename.endswith(".pdf"):
        loader = PyMuPDFLoader(file_path=os.path.join(pdf_dir, filename))
        data.extend(loader.load())

# Ensure there's data, create a demo if not
if len(data) == 0:
    data = [
        Document(
            page_content="Sample text about LummaC2.",
            metadata={"source": "sample.pdf", "page": 1},
        ),
    ]

# Split documents into chunks
data_splits = recur_splitter.split_documents(data)

# (Optional but helpful debug: quickly sanity-check chunks)
# for i, d in enumerate(data_splits[:3]):
#     print(i, len(d.page_content))

# Create ChromaDB vector store (persists locally)
os.makedirs("./chroma_llm_training", exist_ok=True)
vectordb = Chroma.from_documents(
    documents=data_splits,
    embedding=embedding_model,
    collection_name="rag-demo",
    persist_directory="./chroma_llm_training"
)

# Define RAG QA prompt template
qa_template = """
You are a helpful assistant answering questions based on the provided context.
Use the following pieces of context to answer the user's question. If unsure, state you don't know.

Context:
{context}

Question: {question}

Answer:
"""
qa_prompt_template = PromptTemplate.from_template(qa_template)

# Define RetrievalQA chain
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectordb.as_retriever(search_kwargs={"k": 5}),
    return_source_documents=True,
    chain_type="stuff",
    chain_type_kwargs={"prompt": qa_prompt_template},
)

# Test questions
questions = [
    "What is LummaC2?",
    "What is a ClickFix?",
    "Why is the sky blue?"
]

# Run questions through the QA chain
for question in questions:
    print(f"\nQuestion: {question}")
    response = qa_chain({"query": question})
    display(Markdown(response["result"]))

    print("\nSources:")
    seen_sources = set()
    for doc in response["source_documents"][:2]:
        source_info = f"{os.path.basename(doc.metadata['source'])}, Page {doc.metadata.get('page', 'N/A')}"
        if source_info not in seen_sources:
            print(source_info)
            seen_sources.add(source_info)
PDF directory is at: /Users/schwartz/src/genai-essentials/data/pdfs

Question: What is LummaC2?
<IPython.core.display.Markdown object>
Sources:
aa25-141b-threat-actors-deploy-lummac2-malware-to-exfiltrate-sensitive-data-from-organizations.pdf, Page 1
aa25-141b-threat-actors-deploy-lummac2-malware-to-exfiltrate-sensitive-data-from-organizations.pdf, Page 7

Question: What is a ClickFix?
<IPython.core.display.Markdown object>
Sources:
clickfix-attacks-sector-alert-tlpclear.pdf, Page 0

Question: Why is the sky blue?
<IPython.core.display.Markdown object>
Sources:
clickfix-attacks-sector-alert-tlpclear.pdf, Page 1
joint-guidance-enhanced-visibility-hardening-guide-for-comms-infrastructure-508c_0.pdf, Page 0