Iyeoluwa Akinrinola
74 lines
1.9 KiB
Python
74 lines
1.9 KiB
Python
"""
|
|
Test script for Ollama integration.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import requests
|
|
import json
|
|
|
|
# Add the parent directory to the path so we can import ai_service
|
|
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
|
|
|
|
from ai_service.config import config
|
|
from ai_service.models.model_service import model_service
|
|
|
|
def test_available_models():
|
|
"""Test getting available models."""
|
|
models = model_service.get_available_models()
|
|
print("Available models:")
|
|
for model in models:
|
|
print(f"- {model['name']} ({model['id']}): {model['description']}")
|
|
print()
|
|
|
|
def test_generate_response():
|
|
"""Test generating a response."""
|
|
model_id = "llama3.1" # Use a specific model instead of config.DEFAULT_MODEL
|
|
prompt = "What is the capital of France?"
|
|
|
|
print(f"Testing model: {model_id}")
|
|
print(f"Prompt: {prompt}")
|
|
|
|
response = model_service.generate_response(
|
|
model_id=model_id,
|
|
prompt=prompt,
|
|
use_rag=False
|
|
)
|
|
|
|
print("Response:")
|
|
print(response)
|
|
print()
|
|
|
|
def test_rag_response():
|
|
"""Test generating a response with RAG."""
|
|
model_id = "llama3.1" # Use a specific model instead of config.DEFAULT_MODEL
|
|
prompt = "Tell me about the documents in the knowledge base."
|
|
|
|
print(f"Testing RAG with model: {model_id}")
|
|
print(f"Prompt: {prompt}")
|
|
|
|
response = model_service.generate_response(
|
|
model_id=model_id,
|
|
prompt=prompt,
|
|
use_rag=True
|
|
)
|
|
|
|
print("Response with RAG:")
|
|
print(response)
|
|
print()
|
|
|
|
if __name__ == "__main__":
|
|
print("Testing Ollama integration")
|
|
print(f"OpenWebUI URL: {config.OPENWEBUI_URL}")
|
|
|
|
# Override the Ollama API URL to use OpenWebUI
|
|
model_service.ollama_api_url = f"{config.OPENWEBUI_URL}/ollama"
|
|
print(f"Using Ollama API URL: {model_service.ollama_api_url}")
|
|
|
|
print(f"Default model: {config.DEFAULT_MODEL}")
|
|
print()
|
|
|
|
test_available_models()
|
|
test_generate_response()
|
|
test_rag_response()
|