Instead of OpenAI, Can I use llamaIndex and Ollama with Langgraph

Short Course Q&A AI Agents in LangGraph
1

I’m trying to use Ollama and LlamaIndex with LangGraph. The first two lectures are Ok. I can run them without any errors. Here is the sample code:

from llama_index.llms.ollama.base import Ollama
from llama_index.core.llms import ChatMessage
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage, AIMessage
from langgraph.graph import StateGraph, END
from langchain_community.tools.tavily_search import TavilySearchResults
from typing import TypedDict, Annotated
import operator
import re

client = Ollama(
    base_url='http://localhost:11434',
    model="llama3.1",
    temperature=0
)

tool = TavilySearchResults(max_results=4)

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

class Agent:

    def __init__(self, model:Ollama, tools, system=""):
        self.system = system
        graph = StateGraph(AgentState)
        graph.add_node("llm", self.call_llama)
        graph.add_node("action", self.take_action)
        graph.add_conditional_edges(
            "llm",
            self.exists_action,
            {True: "action", False: END}
        )
        graph.add_edge("action", "llm")
        graph.set_entry_point("llm")
        self.graph = graph.compile()
        self.tools = {t.name: t for t in tools}
        self.model = model

    def exists_action(self, state: AgentState):
        result = state['messages'][-1]
        return len(result.tool_calls) > 0
    
    def call_llama(self, state: AgentState):
        messages = []

        for msg in state['messages']:
            msg_dict = msg.model_dump()
            print(msg_dict)
            if msg_dict['type'] == "human":
                role = "user"
            elif msg_dict['type'] == "ai":
                role = "assistant"
            elif msg_dict['type'] == "tool":
                role = "tool"
            else:
                print(msg_dict['type'])
                role = "human"
            messages.append(ChatMessage(role=role, content=msg.content))

        if self.system:
            messages = [ChatMessage(role="system", content=self.system)] + messages
        
        message = self.model.chat(messages=messages)
        message = message.message.content

        return {'messages': [AIMessage(content=message)]}
    
    def take_action(self, state:AgentState):
        tool_calls = state['messages'][-1].tool_calls
        results = []
        for t in tool_calls:
            print(f"Calling: {t}")
            if not t['name'] in self.tool_names:
                print("\n...bad tool name...")
                result = "bad tool name, retry"
            else:
                result = self.tools[t['name']].invoke(t['args'])
            results.append(ToolMessage(tool_call_id=t["id"], name=t['name'], content=str(result)))
        print("Back to the model!")
        return {'messages': results}

prompt = """You are a smart research assistant. Use the search engine to look up information. \
You are allowed to make multiple calls (either together or in sequence). \
Only look up information when you are sure of what you want. \
If you need to look up some information before asking a follow-up question, you are allowed to do that!
"""
abot = Agent(client, [tool], system=prompt)
query = "Who won the super bowl in 2024? In what state is the winning team headquarters located? \
What is the GDP of that state? Answer each question." 
messages = [HumanMessage(content=query)]
result = abot.graph.invoke({"messages": messages})
print(result['messages'][-1].content)

This returns the same answer as the one in lecture 2. However, it doesn’t return or stream the output as in lecture 4. How can I achieve this?

1 Like
2

I found the solution. We can openAI and ChatOpenAI from langchain_openai.

client = OpenAI(
    api_key="ollama",
    base_url="http://localhost:11434/v1"
)

and

from langchain_openai import ChatOpenAI
model = ChatOpenAI(
    model="llama3.1",
    api_key="ollama",
    base_url="http://127.0.0.1:11434/v1"
)

We will need Ollama runs in the backend.

1 Like