wehub-resource-sync
307 lines
11 KiB
Python
307 lines
11 KiB
Python
import os
|
|
from typing import Optional, Any
|
|
import re
|
|
import requests
|
|
|
|
from llama_index.core.workflow import (
|
|
StartEvent,
|
|
StopEvent,
|
|
step,
|
|
Workflow,
|
|
Context,
|
|
)
|
|
from llama_index.core import SummaryIndex
|
|
from llama_index.core.schema import Document
|
|
from llama_index.core.prompts import PromptTemplate
|
|
from llama_index.llms.openai import OpenAI
|
|
from llama_index.core.llms import LLM
|
|
from llama_index.core.base.base_retriever import BaseRetriever
|
|
from typing import List
|
|
|
|
from llama_index.core.schema import NodeWithScore
|
|
from llama_index.core.workflow import (
|
|
Event,
|
|
)
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
|
|
class RetrieveEvent(Event):
|
|
"""Retrieve event (gets retrieved nodes)."""
|
|
|
|
retrieved_nodes: List[NodeWithScore]
|
|
|
|
|
|
class WebSearchEvent(Event):
|
|
"""Web search event."""
|
|
|
|
relevant_text: str # not used, just used for pass through
|
|
|
|
|
|
class QueryEvent(Event):
|
|
"""Query event. Queries given relevant text and search text."""
|
|
|
|
relevant_text: str
|
|
search_text: str
|
|
|
|
|
|
DEFAULT_RELEVANCY_PROMPT_TEMPLATE = PromptTemplate(
|
|
template="""As a grader, your task is to evaluate the relevance of a document retrieved in response to a user's question.
|
|
|
|
Retrieved Document:
|
|
-------------------
|
|
{context_str}
|
|
|
|
User Question:
|
|
--------------
|
|
{query_str}
|
|
|
|
Evaluation Criteria:
|
|
- Consider whether the document contains keywords or topics related to the user's question.
|
|
- The evaluation should not be overly stringent; the primary objective is to identify and filter out clearly irrelevant retrievals.
|
|
|
|
Decision:
|
|
- Assign a binary score to indicate the document's relevance.
|
|
- Use 'yes' if the document is relevant to the question, or 'no' if it is not.
|
|
|
|
Please provide your binary score ('yes' or 'no') below to indicate the document's relevance to the user question."""
|
|
)
|
|
|
|
DEFAULT_TRANSFORM_QUERY_TEMPLATE = PromptTemplate(
|
|
template="""Your task is to refine a query to ensure it is highly effective for retrieving relevant search results. \n
|
|
Analyze the given input to grasp the core semantic intent or meaning. \n
|
|
Original Query:
|
|
\n ------- \n
|
|
{query_str}
|
|
\n ------- \n
|
|
Your goal is to rephrase or enhance this query to improve its search performance. Ensure the revised query is concise and directly aligned with the intended search objective. \n
|
|
Respond with the optimized query only:"""
|
|
)
|
|
|
|
|
|
class CorrectiveRAGWorkflow(Workflow):
|
|
"""Corrective RAG Workflow."""
|
|
|
|
def __init__(
|
|
self,
|
|
index,
|
|
firecrawl_api_key: str,
|
|
llm: Optional[LLM] = None,
|
|
**kwargs: Any
|
|
) -> None:
|
|
"""Init params."""
|
|
super().__init__(**kwargs)
|
|
self.index = index
|
|
self.firecrawl_api_key = firecrawl_api_key
|
|
|
|
if llm is not None:
|
|
self.llm = llm
|
|
else:
|
|
# self.llm = Ollama(
|
|
# model="gemma3:4b",
|
|
# base_url="http://localhost:11434",
|
|
# temperature=0.1,
|
|
# )
|
|
self.llm = OpenAI(
|
|
model="gpt-4o",
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
)
|
|
|
|
# Set the global LLM settings to avoid conflicts
|
|
from llama_index.core import Settings
|
|
Settings.llm = self.llm
|
|
|
|
@step
|
|
async def retrieve(self, ctx: Context, ev: StartEvent) -> Optional[RetrieveEvent]:
|
|
"""Retrieve the relevant nodes for the query."""
|
|
query_str = ev.get("query_str")
|
|
retriever_kwargs = ev.get("retriever_kwargs", {})
|
|
|
|
print(f"DEBUG: retrieve - query_str: {query_str}")
|
|
print(f"DEBUG: retrieve - retriever_kwargs: {retriever_kwargs}")
|
|
|
|
if query_str is None:
|
|
print("DEBUG: retrieve - query_str is None, returning None")
|
|
return None
|
|
|
|
retriever: BaseRetriever = self.index.as_retriever(**retriever_kwargs)
|
|
print(f"DEBUG: retrieve - retriever created: {type(retriever)}")
|
|
|
|
result = retriever.retrieve(query_str)
|
|
print(f"DEBUG: retrieve - retrieved {len(result)} nodes")
|
|
|
|
if result:
|
|
print(f"DEBUG: retrieve - first node preview: {result[0].text[:100]}...")
|
|
|
|
await ctx.set("retrieved_nodes", result)
|
|
await ctx.set("query_str", query_str)
|
|
return RetrieveEvent(retrieved_nodes=result)
|
|
|
|
@step
|
|
async def eval_relevance(
|
|
self, ctx: Context, ev: RetrieveEvent
|
|
) -> WebSearchEvent | QueryEvent:
|
|
"""Evaluate relevancy of retrieved documents with the query."""
|
|
retrieved_nodes = ev.retrieved_nodes
|
|
query_str = await ctx.get("query_str")
|
|
|
|
print(f"DEBUG: Retrieved {len(retrieved_nodes)} nodes")
|
|
print(f"DEBUG: Query: {query_str}")
|
|
|
|
relevancy_results = []
|
|
for i, node in enumerate(retrieved_nodes):
|
|
print(f"DEBUG: Node {i} text preview: {node.text[:100]}...")
|
|
prompt = DEFAULT_RELEVANCY_PROMPT_TEMPLATE.format(
|
|
context_str=node.text, query_str=query_str)
|
|
try:
|
|
relevancy = await self.llm.acomplete(prompt)
|
|
relevancy_results.append(relevancy.text.lower().strip())
|
|
print(f"DEBUG: Node {i} relevancy: {relevancy.text}")
|
|
except Exception as e:
|
|
# Fallback to synchronous call if async is not supported
|
|
relevancy = self.llm.complete(prompt)
|
|
relevancy_results.append(relevancy.text.lower().strip())
|
|
print(f"DEBUG: Node {i} relevancy (sync): {relevancy.text}")
|
|
|
|
print(f"DEBUG: All relevancy results: {relevancy_results}")
|
|
|
|
relevancy_results_striped = [re.sub(r"<think>.*?</think>", "", s, flags=re.DOTALL).strip() for s in relevancy_results]
|
|
|
|
# Improved relevancy parsing - look for "yes" anywhere in the response
|
|
relevant_texts = [
|
|
retrieved_nodes[i].text
|
|
for i, result in enumerate(relevancy_results_striped)
|
|
if "yes" in result.lower()
|
|
]
|
|
relevant_text = "\n".join(relevant_texts)
|
|
|
|
print(f"DEBUG: Relevant texts count: {len(relevant_texts)}")
|
|
print(f"DEBUG: Relevant text preview: {relevant_text[:200]}...")
|
|
|
|
if "no" in relevancy_results_striped:
|
|
print("DEBUG: Some documents irrelevant, returning WebSearchEvent")
|
|
return WebSearchEvent(relevant_text=relevant_text)
|
|
else:
|
|
print("DEBUG: All documents relevant, returning QueryEvent")
|
|
return QueryEvent(relevant_text=relevant_text, search_text="")
|
|
|
|
def _firecrawl_search(self, query: str, limit: int = 5) -> str:
|
|
"""Perform web search using FireCrawl API directly."""
|
|
url = "https://api.firecrawl.dev/v1/search"
|
|
|
|
payload = {
|
|
"query": query,
|
|
"limit": 5,
|
|
"timeout": 60000,
|
|
}
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {self.firecrawl_api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, json=payload, headers=headers, timeout=60)
|
|
response.raise_for_status()
|
|
|
|
data = response.json()
|
|
|
|
if data.get("success") and data.get("data"):
|
|
# Extract title and description from each result
|
|
search_results = []
|
|
for result in data["data"]:
|
|
title = result.get("title", "")
|
|
description = result.get("description", "")
|
|
url = result.get("url", "")
|
|
|
|
if title or description:
|
|
result_text = f"Title: {title}\nDescription: {description}\nURL: {url}\n"
|
|
search_results.append(result_text)
|
|
|
|
return "\n---\n".join(search_results)
|
|
else:
|
|
print(f"DEBUG: FireCrawl API returned no results or error: {data}")
|
|
return ""
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"DEBUG: FireCrawl API request failed: {e}")
|
|
return ""
|
|
except Exception as e:
|
|
print(f"DEBUG: Unexpected error in FireCrawl search: {e}")
|
|
return ""
|
|
|
|
@step
|
|
async def web_search(
|
|
self, ctx: Context, ev: WebSearchEvent
|
|
) -> QueryEvent:
|
|
"""Search the transformed query"""
|
|
# If any document is found irrelevant, transform the query string for better search results.
|
|
|
|
query_str = await ctx.get("query_str")
|
|
|
|
prompt = DEFAULT_TRANSFORM_QUERY_TEMPLATE.format(query_str=query_str)
|
|
try:
|
|
result = await self.llm.acomplete(prompt)
|
|
transformed_query_str = result.text
|
|
except Exception as e:
|
|
# Fallback to synchronous call if async is not supported
|
|
result = self.llm.complete(prompt)
|
|
transformed_query_str = result.text
|
|
|
|
print(f"DEBUG: web_search - transformed query: {transformed_query_str}")
|
|
|
|
# Conduct a search with the transformed query string using direct API call
|
|
search_text = self._firecrawl_search(transformed_query_str)
|
|
|
|
print(f"DEBUG: web_search - search results length: {len(search_text)}")
|
|
if search_text:
|
|
print(f"DEBUG: web_search - search results preview: {search_text[:200]}...")
|
|
|
|
return QueryEvent(relevant_text=ev.relevant_text, search_text=search_text)
|
|
|
|
@step
|
|
async def query_result(self, ctx: Context, ev: QueryEvent) -> StopEvent:
|
|
"""Get result with relevant text."""
|
|
relevant_text = ev.relevant_text
|
|
search_text = ev.search_text
|
|
query_str = await ctx.get("query_str")
|
|
|
|
print(f"DEBUG: query_result - query_str: {query_str}")
|
|
print(f"DEBUG: query_result - relevant_text: {relevant_text}")
|
|
print(f"DEBUG: query_result - search_text: {search_text}")
|
|
|
|
if not relevant_text.strip() and not search_text.strip():
|
|
print("DEBUG: No relevant text, returning empty response")
|
|
return StopEvent(result="No relevant information found in the documents.")
|
|
|
|
|
|
context_str = relevant_text + "\n" + search_text
|
|
|
|
prompt = f"""As a helpful assistant, your task is to answer the user's question based on the given context.
|
|
|
|
A few things to keep in mind:
|
|
- The context can either be relevant text or web search results.
|
|
- The context can also be a mix of both.
|
|
|
|
Your task is to look at the query and the whole context and generate what you think is the best answer to the question.
|
|
|
|
Here is the context:
|
|
Context:
|
|
{context_str}
|
|
|
|
--------------------------------
|
|
|
|
Question:
|
|
{query_str}
|
|
|
|
--------------------------------
|
|
|
|
Generate an answer to the question:
|
|
"""
|
|
|
|
result = await self.llm.acomplete(prompt)
|
|
print(f"DEBUG: query_result - final result: {result.text}")
|
|
return StopEvent(result=result.text)
|