Ollama + LangGraph 웹 검색 챗봇

Google Colab에서 Ollama를 사용한 무료 웹 검색 챗봇을 구현합니다.

1. 필요한 패키지 설치

# 필요한 패키지 설치
!pip install -qU tavily-python langchain_community langgraph langchain_core pydantic
!pip install -qU requests aiohttp

2. Colab에서 Ollama 설치 및 설정

# Ollama 설치
!curl -fsSL https://ollama.ai/install.sh | sh

# Ollama 서버를 백그라운드에서 시작
import subprocess
import time
import threading

# Ollama 서버 시작
def start_ollama_server():
    try:
        subprocess.run(['/usr/local/bin/ollama', 'serve'], check=False)
    except:
        pass

# 백그라운드에서 Ollama 서버 실행
server_thread = threading.Thread(target=start_ollama_server, daemon=True)
server_thread.start()

# 서버가 시작될 때까지 잠시 대기
print("Ollama 서버 시작 중...")
time.sleep(5)
print("Ollama 서버 시작 완료!")

# 경량 모델 다운로드 (처음 실행 시에만 필요)
import subprocess
import time

print("모델 다운로드 중... (몇 분 소요될 수 있습니다)")
try:
    # 작은 모델부터 시도
    result = subprocess.run(['/usr/local/bin/ollama', 'pull', 'llama3.2:1b'], 
                          capture_output=True, text=True, timeout=600)
    if result.returncode == 0:
        print("llama3.2:1b 모델 다운로드 완료!")
        model_name = "llama3.2:1b"
    else:
        print("llama3.2:1b 실패, 대체 모델 시도 중...")
        result = subprocess.run(['/usr/local/bin/ollama', 'pull', 'tinyllama'], 
                              capture_output=True, text=True, timeout=600)
        if result.returncode == 0:
            print("tinyllama 모델 다운로드 완료!")
            model_name = "tinyllama"
        else:
            raise Exception("모델 다운로드 실패")
except Exception as e:
    print(f"오류 발생: {e}")
    print("수동으로 모델을 설정합니다...")
    model_name = "tinyllama"

print(f"사용할 모델: {model_name}")

3. API 키 설정

import os

# Tavily API 키 설정 (https://app.tavily.com/home 에서 발급)
os.environ["TAVILY_API_KEY"] = "tvly-YOUR_API_KEY_HERE"  # 여기에 실제 API 키 입력

# API 키가 설정되었는지 확인
if os.environ["TAVILY_API_KEY"] == "tvly-YOUR_API_KEY_HERE":
    print("Tavily API 키를 설정해주세요!")
else:
    print("API 키 설정 완료!")

4. 커스텀 Ollama 클라이언트 구현 (JSON 에러 해결)

import requests
import json
from typing import List, Dict, Any, Optional
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage, AIMessage, HumanMessage, SystemMessage
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from pydantic import Field

class CustomOllamaChat(BaseChatModel):
    """JSON 에러를 해결한 커스텀 Ollama 클라이언트"""
    
    model_name: str = Field(default="tinyllama")
    base_url: str = Field(default="http://localhost:11434")
    temperature: float = Field(default=0.7)
    
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ):
        # 메시지를 Ollama 형식으로 변환
        prompt = self._messages_to_prompt(messages)
        
        # Ollama API 호출
        try:
            response = requests.post(
                f"{self.base_url}/api/generate",
                json={
                    "model": self.model_name,
                    "prompt": prompt,
                    "stream": False,
                    "options": {
                        "temperature": self.temperature,
                        "num_predict": 150
                    }
                },
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                content = result.get("response", "응답을 생성할 수 없습니다.")
            else:
                content = f"API 호출 실패 (상태 코드: {response.status_code})"
                
        except requests.exceptions.RequestException as e:
            content = f"연결 오류: {str(e)}"
        except json.JSONDecodeError as e:
            content = f"JSON 디코딩 오류: {str(e)}"
        except Exception as e:
            content = f"예상치 못한 오류: {str(e)}"
        
        # ChatResult 형태로 반환 (간단한 형태)
        from langchain_core.outputs import ChatResult, ChatGeneration
        message = AIMessage(content=content)
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])
    
    def _messages_to_prompt(self, messages: List[BaseMessage]) -> str:
        """메시지 리스트를 단일 프롬프트 문자열로 변환"""
        prompt_parts = []
        
        for message in messages:
            if isinstance(message, HumanMessage):
                prompt_parts.append(f"Human: {message.content}")
            elif isinstance(message, AIMessage):
                prompt_parts.append(f"Assistant: {message.content}")
            elif isinstance(message, SystemMessage):
                prompt_parts.append(f"System: {message.content}")
            else:
                prompt_parts.append(f"{message.content}")
        
        return "\n".join(prompt_parts) + "\nAssistant:"
    
    @property
    def _llm_type(self) -> str:
        return "custom_ollama"
    
    def bind_tools(self, tools):
        """도구 바인딩 (단순 구현)"""
        return CustomOllamaWithTools(model_name=self.model_name, 
                                   base_url=self.base_url, 
                                   temperature=self.temperature, 
                                   tools=tools)

class CustomOllamaWithTools(CustomOllamaChat):
    """도구를 지원하는 Ollama 클라이언트"""
    
    tools: List = Field(default_factory=list)
    
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ):
        # 사용자 메시지에서 검색 키워드 감지
        user_message = ""
        for msg in messages:
            if isinstance(msg, HumanMessage):
                user_message = msg.content
                break
        
        # 검색이 필요한지 판단
        search_keywords = ["검색", "찾아", "알려줘", "what", "how", "when", "where", "최신", "뉴스", "정보"]
        needs_search = any(keyword.lower() in user_message.lower() for keyword in search_keywords)
        
        if needs_search and self.tools:
            # 검색 수행
            search_tool = self.tools[0]  # 첫 번째 도구가 검색 도구라고 가정
            try:
                search_results = search_tool.invoke({"query": user_message})
                
                # 검색 결과를 포함한 새로운 프롬프트 생성
                enhanced_prompt = f"""
다음 검색 결과를 바탕으로 질문에 답해주세요:

검색 결과: {json.dumps(search_results, ensure_ascii=False)}

질문: {user_message}

답변:"""
                
                # 검색 결과를 포함한 메시지로 LLM 호출
                enhanced_messages = [HumanMessage(content=enhanced_prompt)]
                return super()._generate(enhanced_messages, stop, run_manager, **kwargs)
                
            except Exception as e:
                print(f"검색 중 오류: {e}")
        
        # 일반 응답
        return super()._generate(messages, stop, run_manager, **kwargs)

print("커스텀 Ollama 클라이언트 준비 완료!")

5. LangGraph 챗봇 설정

from typing import Annotated
from typing_extensions import TypedDict
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# State 정의
class State(TypedDict):
    messages: Annotated[list, add_messages]

# 모델 초기화
llm = CustomOllamaChat(model_name=model_name)

# 검색 도구 설정
tool = TavilySearchResults(max_results=2)
tools = [tool]
llm_with_tools = llm.bind_tools(tools)

# 그래프 빌더 생성
graph_builder = StateGraph(State)

print("LangGraph 설정 완료!")

# 챗봇 노드 정의
def chatbot(state: State):
    try:
        response = llm_with_tools.invoke(state["messages"])
        return {"messages": [response]}
    except Exception as e:
        error_message = f"응답 생성 중 오류가 발생했습니다: {str(e)}"
        from langchain_core.messages import AIMessage
        return {"messages": [AIMessage(content=error_message)]}

# 노드 추가
graph_builder.add_node("chatbot", chatbot)

# 엣지 연결 (단순화된 버전)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)

# 그래프 컴파일
graph = graph_builder.compile()

print("챗봇 그래프 구성 완료!")

6. 챗봇 테스트

# 스트리밍 함수
def stream_graph_updates(user_input: str):
    try:
        from langchain_core.messages import HumanMessage
        
        print(f"User: {user_input}")
        
        for event in graph.stream({"messages": [HumanMessage(content=user_input)]}):
            for value in event.values():
                if value["messages"]:
                    message = value["messages"][-1]
                    if hasattr(message, 'content'):
                        print(f"Assistant: {message.content}")
                    else:
                        print(f"Assistant: {message}")
    except Exception as e:
        print(f"오류 발생: {e}")

# 간단한 테스트
print("=== 챗봇 테스트 ===")
stream_graph_updates("안녕하세요!")

# 검색 기능 테스트
print("\n=== 검색 기능 테스트 ===")
stream_graph_updates("최신 AI 뉴스를 알려줘")

7. 대화형 챗봇 실행

# 대화형 챗봇
def run_interactive_chatbot():
    print("Ollama 기반 웹 검색 챗봇입니다!")
    print("종료하려면 'quit', 'exit', 'q' 중 하나를 입력하세요.\n")
    
    conversation_history = []
    
    while True:
        try:
            user_input = input("User: ")
            
            if user_input.lower() in ["quit", "exit", "q", "종료"]:
                print("안녕히 가세요!")
                break
            
            if not user_input.strip():
                continue
            
            # 대화 기록에 추가
            from langchain_core.messages import HumanMessage
            conversation_history.append(HumanMessage(content=user_input))
            
            # 응답 생성
            try:
                for event in graph.stream({"messages": conversation_history}):
                    for value in event.values():
                        if value["messages"]:
                            message = value["messages"][-1]
                            response_content = message.content if hasattr(message, 'content') else str(message)
                            print(f"Assistant: {response_content}")
                            
                            # 응답을 대화 기록에 추가
                            from langchain_core.messages import AIMessage
                            conversation_history.append(AIMessage(content=response_content))
            except Exception as e:
                print(f"응답 생성 중 오류: {e}")
            
            print()  # 빈 줄 추가
            
        except KeyboardInterrupt:
            print("\n안녕히 가세요!")
            break
        except Exception as e:
            print(f"오류 발생: {e}")

# 대화형 챗봇 실행
run_interactive_chatbot()

8. 트러블슈팅 및 유용한 명령어

# Ollama 상태 확인
import subprocess
import requests

def check_ollama_status():
    print("=== Ollama 상태 확인 ===")
    
    # 1. Ollama 서버 상태 확인
    try:
        response = requests.get("http://localhost:11434/api/tags", timeout=5)
        if response.status_code == 200:
            print("Ollama 서버가 실행 중입니다.")
            models = response.json()
            if models.get('models'):
                print(f"설치된 모델: {[model['name'] for model in models['models']]}")
            else:
                print("설치된 모델이 없습니다.")
        else:
            print(f"Ollama 서버 응답 오류 (상태 코드: {response.status_code})")
    except requests.exceptions.ConnectionError:
        print("Ollama 서버에 연결할 수 없습니다. 서버가 실행 중인지 확인하세요.")
    except Exception as e:
        print(f"오류: {e}")
    
    # 2. Tavily API 키 확인
    if os.environ.get("TAVILY_API_KEY", "").startswith("tvly-"):
        print("Tavily API 키가 설정되어 있습니다.")
    else:
        print("Tavily API 키를 확인하세요.")

check_ollama_status()

# 수동으로 Ollama 서버 재시작
def restart_ollama():
    print("Ollama 서버 재시작 중...")
    
    # 기존 서버 종료 시도
    try:
        subprocess.run(["pkill", "-f", "ollama"], check=False)
        time.sleep(2)
    except:
        pass
    
    # 새 서버 시작
    server_thread = threading.Thread(target=start_ollama_server, daemon=True)
    server_thread.start()
    time.sleep(5)
    
    check_ollama_status()

# 필요시 주석 해제 후 실행
# restart_ollama()