Refactor OllamaAssist to use LangChain

[madtank]

Refactor OllamaAssist to use LangChain

Objective

Refactor the OllamaAssist project to integrate LangChain, improving code structure, flexibility, and enabling advanced features like dynamic tool usage and improved context management.

Tasks

1. Setup and Preparation

• [ ] Install LangChain and dependencies:

  pip install langchain langchain-community langchain-openai

• [ ] Update requirements.txt with new dependencies

2. Refactor llm_helper.py

• [ ] Add LangChain imports:

  from langchain_community.chat_models import ChatOllama
  from langchain.schema import HumanMessage, AIMessage, SystemMessage
  from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
  from langchain.callbacks import AsyncIteratorCallbackHandler

• [ ] Replace chat function with LangChain implementation:

  def chat(messages, model, tools=None, stream=True):
    # Convert messages to LangChain format
    langchain_messages = [
        SystemMessage(content=msg["content"]) if msg["role"] == "system" else
        HumanMessage(content=msg["content"]) if msg["role"] == "user" else
        AIMessage(content=msg["content"]) if msg["role"] == "assistant" else
        HumanMessage(content=f"Function {msg.get('name', 'unknown')} returned: {msg['content']}")
        for msg in messages
    ]
  
    # Create ChatOllama instance
    chat_model = ChatOllama(
        model=model,
        streaming=stream,
        callbacks=[StreamingStdOutCallbackHandler()] if stream else None
    )
  
    if stream:
        callback = AsyncIteratorCallbackHandler()
        chat_model = ChatOllama(
            model=model,
            streaming=True,
            callbacks=[callback]
        )
  
        async def async_generator():
            async for chunk in callback.aiter():
                yield {
                    "message": {
                        "role": "assistant",
                        "content": chunk.content
                    }
                }
  
        return async_generator()
    else:
        response = chat_model(langchain_messages)
        return {
            "message": {
                "role": "assistant",
                "content": response.content
            }
        }

3. Refactor tools.py

• [ ] Convert tools to LangChain format:

  from langchain.tools import Tool
  from langchain.pydantic_v1 import BaseModel, Field
  
  class SearchInput(BaseModel):
    query: str = Field(..., description="The search query")
  
  search_tool = Tool(
    name="search_duckduckgo",
    description="DuckDuckGo web search. You write the query, and we will get the results back directly in our chat.",
    func=search_duckduckgo,
    args_schema=SearchInput
  )
  
  class BreakthroughInput(BaseModel):
    internal_dialogue: str = Field(..., description="Your comprehensive internal analysis, planning, and problem-solving thoughts. This is not a message to the user.")
  
  breakthrough_tool = Tool(
    name="breakthrough_blast",
    description="Internal thought process for deep analysis and planning. Use this for private reflection, not direct user communication.",
    func=breakthrough_blast,
    args_schema=BreakthroughInput
  )
  
  tools = [search_tool, breakthrough_tool]

• [ ] Remove tools.yaml file

4. Refactor ollama_chatbot.py

• [ ] Update imports:

  from langchain.agents import AgentExecutor, create_openai_tools_agent
  from langchain.prompts import ChatPromptTemplate
  from langchain.callbacks import StreamlitCallbackHandler

• [ ] Modify generate_response function:

  from src.tools import tools
  
  def generate_response(model, use_tools):
    if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
        with st.spinner('Generating response...'):
            stream_handler = StreamlitCallbackHandler(st.empty())
            chat_model = ChatOllama(model=model, streaming=True, callbacks=[stream_handler])
  
            if use_tools:
                prompt = ChatPromptTemplate.from_messages([
                    ("system", "You are Luna, an advanced AI assistant..."),  # Your system prompt
                    ("human", "{input}")
                ])
                agent = create_openai_tools_agent(chat_model, tools, prompt)
                agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
  
                with st.chat_message("assistant"):
                    response = agent_executor.invoke({
                        "input": st.session_state.messages[-1]["content"]
                    }, callbacks=[stream_handler])
  
                assistant_response = response['output']
            else:
                messages = [SystemMessage(content="You are Luna, an advanced AI assistant...")] + \
                           [HumanMessage(content=m["content"]) if m["role"] == "user" else 
                            AIMessage(content=m["content"]) for m in st.session_state.messages]
  
                with st.chat_message("assistant"):
                    response = chat_model(messages)
  
                assistant_response = response.content
  
            st.session_state.messages.append({"role": "assistant", "content": assistant_response})

5. Testing and Debugging

• [ ] Test basic chat functionality without tools
• [ ] Test each tool individually
• [ ] Test streaming responses
• [ ] Verify agent tool selection and usage

6. Further Enhancements (Optional)

• [ ] Implement LangChain's memory systems
• [ ] Explore advanced agent types
• [ ] Implement dynamic prompt templates

Notes

• Pay attention to Ollama-specific features and ensure compatibility with LangChain
• Monitor performance, especially for streaming responses
• Be prepared to adjust the agent type or create a custom agent if needed for optimal Ollama integration

Resources

• LangChain Documentation
• Ollama GitHub Repository