mirror of
https://github.com/volcengine/OpenViking.git
synced 2026-09-30 17:28:07 +08:00
Qin Haojie
7abd6ab249
* refactor(client): remove Python embedded mode Consolidate Python consumers on the HTTP SDK while keeping shared server and storage capabilities unchanged. * refactor(client): remove obsolete embedded leftovers
252 lines
8.5 KiB
Python
252 lines
8.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
RAG Pipeline - Retrieval-Augmented Generation using OpenViking + LLM
|
|
Focused on querying and answer generation, not resource management
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import requests
|
|
from openviking_sdk import SyncHTTPClient
|
|
|
|
|
|
class Recipe:
|
|
"""
|
|
Recipe (Boring name is RAG Pipeline)
|
|
|
|
Combines semantic search with LLM generation:
|
|
1. Search OpenViking database for relevant context
|
|
2. Send context + query to LLM
|
|
3. Return generated answer with sources
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config_path: str = "./ov.conf",
|
|
server_url: str = "http://127.0.0.1:1933",
|
|
):
|
|
"""
|
|
Initialize RAG pipeline
|
|
|
|
Args:
|
|
config_path: Path to config file with LLM settings
|
|
server_url: OpenViking HTTP server URL
|
|
"""
|
|
# Load configuration
|
|
with open(config_path, "r") as f:
|
|
self.config_dict = json.load(f)
|
|
|
|
# Extract LLM config
|
|
self.vlm_config = self.config_dict.get("vlm", {})
|
|
self.api_base = self.vlm_config.get("api_base")
|
|
self.api_key = self.vlm_config.get("api_key")
|
|
self.model = self.vlm_config.get("model")
|
|
|
|
# Initialize OpenViking client
|
|
self.client = SyncHTTPClient(url=server_url)
|
|
self.client.initialize()
|
|
|
|
def search(
|
|
self,
|
|
query: str,
|
|
top_k: int = 3,
|
|
target_uri: Optional[str] = None,
|
|
score_threshold: float = 0.2,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search for relevant content using semantic search
|
|
|
|
Args:
|
|
query: Search query
|
|
top_k: Number of results to return
|
|
target_uri: Optional specific URI to search in. If None, searches all resources.
|
|
score_threshold: Minimum relevance score for search results (default: 0.2)
|
|
|
|
Returns:
|
|
List of search results with content and scores
|
|
"""
|
|
# print(f"🔍 Searching for: '{query}'")
|
|
|
|
# Search all resources or specific target
|
|
# `find` has better performance, but not so smart
|
|
results = self.client.search(query, target_uri=target_uri, score_threshold=score_threshold)
|
|
|
|
# Extract top results
|
|
search_results = []
|
|
resources = results.get("resources", [])[:top_k]
|
|
memories = results.get("memories", [])[:top_k]
|
|
for _i, resource in enumerate(resources + memories): # ignore SKILLs for mvp
|
|
uri = resource["uri"]
|
|
score = resource.get("score", 0.0)
|
|
try:
|
|
content = self.client.read(uri)
|
|
search_results.append(
|
|
{
|
|
"uri": uri,
|
|
"score": score,
|
|
"content": content,
|
|
}
|
|
)
|
|
# print(f" {i + 1}. {uri} (score: {score:.4f})")
|
|
|
|
except Exception as e:
|
|
# Handle directories - read their abstract instead
|
|
if "is a directory" in str(e):
|
|
try:
|
|
abstract = self.client.abstract(uri)
|
|
search_results.append(
|
|
{
|
|
"uri": uri,
|
|
"score": score,
|
|
"content": f"[Directory Abstract] {abstract}",
|
|
}
|
|
)
|
|
# print(f" {i + 1}. {uri} (score: {score:.4f}) [directory]")
|
|
except:
|
|
# Skip if we can't get abstract
|
|
continue
|
|
else:
|
|
# Skip other errors
|
|
continue
|
|
|
|
return search_results
|
|
|
|
def call_llm(
|
|
self, messages: List[Dict[str, str]], temperature: float = 0.7, max_tokens: int = 2048
|
|
) -> str:
|
|
"""
|
|
Call LLM API to generate response
|
|
|
|
Args:
|
|
messages: List of message dictionaries with 'role' and 'content' keys
|
|
Each message should have format: {"role": "user|assistant|system", "content": "..."}
|
|
temperature: Sampling temperature (0.0 to 1.0)
|
|
max_tokens: Maximum tokens to generate
|
|
|
|
Returns:
|
|
LLM response text
|
|
"""
|
|
url = f"{self.api_base}/chat/completions"
|
|
|
|
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}
|
|
|
|
payload = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
|
|
print(f"🤖 Calling LLM: {self.model}")
|
|
response = requests.post(url, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
answer = result["choices"][0]["message"]["content"]
|
|
|
|
return answer
|
|
|
|
def query(
|
|
self,
|
|
user_query: str,
|
|
search_top_k: int = 3,
|
|
temperature: float = 0.7,
|
|
max_tokens: int = 2048,
|
|
system_prompt: Optional[str] = None,
|
|
score_threshold: float = 0.2,
|
|
chat_history: Optional[List[Dict[str, str]]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Full RAG pipeline: search → retrieve → generate
|
|
|
|
Args:
|
|
user_query: User's question
|
|
search_top_k: Number of search results to use as context
|
|
temperature: LLM sampling temperature
|
|
max_tokens: Maximum tokens to generate
|
|
system_prompt: Optional system prompt to prepend
|
|
score_threshold: Minimum relevance score for search results (default: 0.2)
|
|
chat_history: Optional list of previous conversation turns for multi-round chat.
|
|
Each turn should be a dict with 'role' and 'content' keys.
|
|
Example: [{"role": "user", "content": "previous question"},
|
|
{"role": "assistant", "content": "previous answer"}]
|
|
|
|
Returns:
|
|
Dictionary with answer, context, metadata, and timings
|
|
"""
|
|
# Track total time
|
|
start_total = time.perf_counter()
|
|
|
|
# Step 1: Search for relevant content (timed)
|
|
start_search = time.perf_counter()
|
|
search_results = self.search(
|
|
user_query, top_k=search_top_k, score_threshold=score_threshold
|
|
)
|
|
search_time = time.perf_counter() - start_search
|
|
|
|
# Step 2: Build context from search results
|
|
context_text = "no relevant information found, try answer based on existing knowledge."
|
|
if search_results:
|
|
context_text = (
|
|
"Answer should pivoting to the following:\n<context>\n"
|
|
+ "\n\n".join(
|
|
[
|
|
f"[Source {i + 1}] (relevance: {r['score']:.4f})\n{r['content']}"
|
|
for i, r in enumerate(search_results)
|
|
]
|
|
)
|
|
+ "\n</context>"
|
|
)
|
|
|
|
# Step 3: Build messages array for chat completion API
|
|
messages = []
|
|
|
|
# Add system message if provided
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
else:
|
|
messages.append(
|
|
{
|
|
"role": "system",
|
|
"content": "Answer questions with plain text. avoid markdown special character",
|
|
}
|
|
)
|
|
|
|
# Add chat history if provided (for multi-round conversations)
|
|
if chat_history:
|
|
messages.extend(chat_history)
|
|
|
|
# Build current turn prompt with context and question
|
|
current_prompt = f"{context_text}\n"
|
|
current_prompt += f"Question: {user_query}\n\n"
|
|
|
|
# Add current user message
|
|
messages.append({"role": "user", "content": current_prompt})
|
|
|
|
# Step 4: Call LLM with messages array (timed)
|
|
start_llm = time.perf_counter()
|
|
answer = self.call_llm(messages, temperature=temperature, max_tokens=max_tokens)
|
|
llm_time = time.perf_counter() - start_llm
|
|
|
|
# Calculate total time
|
|
total_time = time.perf_counter() - start_total
|
|
|
|
# Return full result with timing data
|
|
return {
|
|
"answer": answer,
|
|
"context": search_results,
|
|
"query": user_query,
|
|
"prompt": current_prompt,
|
|
"timings": {
|
|
"search_time": search_time,
|
|
"llm_time": llm_time,
|
|
"total_time": total_time,
|
|
},
|
|
}
|
|
|
|
def close(self):
|
|
"""Clean up resources"""
|
|
self.client.close()
|