mirrors/patchy631--ai-engineering-hub
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:37:47 +08:00
commit 7653f56fed
1422 changed files with 359026 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
corrective-rag/app.py
+235
View File
@@ -0,0 +1,235 @@
import nest_asyncio
nest_asyncio.apply()
from dotenv import load_dotenv
load_dotenv()
import logging
import sys
import os
import asyncio
import streamlit as st
import qdrant_client
import base64
import gc
import tempfile
import uuid
import time
from IPython.display import Markdown, display
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core import StorageContext
from llama_index.llms.ollama import Ollama
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.embeddings.fastembed import FastEmbedEmbedding
from llama_index.core import Settings
from workflow import CorrectiveRAGWorkflow
import io
from contextlib import redirect_stdout
# Set up page configuration
st.set_page_config(page_title="Corrective RAG Demo", layout="wide")
# Initialize session state variables
if "id" not in st.session_state:
st.session_state.id = uuid.uuid4()
st.session_state.file_cache = {}
if "workflow" not in st.session_state:
st.session_state.workflow = None
if "messages" not in st.session_state:
st.session_state.messages = []
if "workflow_logs" not in st.session_state:
st.session_state.workflow_logs = []
session_id = st.session_state.id
@st.cache_resource
def load_llm():
llm = Ollama(model="deepseek-r1:7b", request_timeout=120.0)
return llm
def reset_chat():
st.session_state.messages = []
gc.collect()
def display_pdf(file):
st.markdown("### PDF Preview")
base64_pdf = base64.b64encode(file.read()).decode("utf-8")
# Embedding PDF in HTML
pdf_display = f"""<iframe src="data:application/pdf;base64,{base64_pdf}" width="400" height="100%" type="application/pdf"
style="height:100vh; width:100%"
>
</iframe>"""
# Displaying File
st.markdown(pdf_display, unsafe_allow_html=True)
# Function to initialize the workflow with uploaded documents
def initialize_workflow(file_path):
with st.spinner("Loading documents and initializing the workflow..."):
documents = SimpleDirectoryReader(file_path).load_data()
client = qdrant_client.QdrantClient(
host="localhost",
port=6333
)
vector_store = QdrantVectorStore(client=client, collection_name="test")
embed_model = FastEmbedEmbedding(model_name="BAAI/bge-large-en-v1.5")
Settings.embed_model = embed_model
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
)
workflow = CorrectiveRAGWorkflow(
index=index,
linkup_api_key=os.environ["LINKUP_API_KEY"],
verbose=True,
timeout=60,
llm=load_llm()
)
st.session_state.workflow = workflow
return workflow
# Function to run the async workflow
async def run_workflow(query):
# Capture stdout to get the workflow logs
f = io.StringIO()
with redirect_stdout(f):
result = await st.session_state.workflow.run(query_str=query)
# Get the captured logs and store them
logs = f.getvalue()
if logs:
st.session_state.workflow_logs.append(logs)
return result
# Sidebar for document upload
with st.sidebar:
# Add Linkup logo and Configuration header in the same line
col1, col2 = st.columns([1, 3])
with col1:
# Add vertical space to align with header
st.write("")
st.image("./assets/linkup.png", width=65)
with col2:
st.header("Linkup Configuration")
st.write("Deep Web Search")
# Add hyperlink to get API key
st.markdown("[Get your API key](https://app.linkup.so/sign-up)", unsafe_allow_html=True)
linkup_api_key = st.text_input("Enter your Linkup API Key", type="password")
# Store API key as environment variable
if linkup_api_key:
os.environ["LINKUP_API_KEY"] = linkup_api_key
st.success("API Key stored successfully!")
st.header("Add your documents!")
uploaded_file = st.file_uploader("Choose your `.pdf` file", type="pdf")
if uploaded_file:
try:
with tempfile.TemporaryDirectory() as temp_dir:
file_path = os.path.join(temp_dir, uploaded_file.name)
with open(file_path, "wb") as f:
f.write(uploaded_file.getvalue())
file_key = f"{session_id}-{uploaded_file.name}"
st.write("Indexing your document...")
if file_key not in st.session_state.get('file_cache', {}):
# Initialize workflow with the uploaded document
workflow = initialize_workflow(temp_dir)
st.session_state.file_cache[file_key] = workflow
else:
st.session_state.workflow = st.session_state.file_cache[file_key]
# Inform the user that the file is processed and Display the PDF uploaded
st.success("Ready to Chat!")
display_pdf(uploaded_file)
except Exception as e:
st.error(f"An error occurred: {e}")
st.stop()
# Main chat interface
col1, col2 = st.columns([6, 1])
with col1:
# Removed the original header
st.markdown("<h2 style='color: #0066cc;'>⚙️ Corrective RAG agentic workflow</h2>", unsafe_allow_html=True)
# Replace text with image and subtitle styling
st.markdown("<div style='display: flex; align-items: center; gap: 10px;'><span style='font-size: 28px; color: #666;'>Powered by LlamaIndex</span><img src='data:image/png;base64,{}' width='50'></div>".format(
base64.b64encode(open("./assets/llamaindex.png", "rb").read()).decode()
), unsafe_allow_html=True)
with col2:
st.button("Clear ↺", on_click=reset_chat)
# Display chat messages from history on app rerun
for i, message in enumerate(st.session_state.messages):
with st.chat_message(message["role"]):
st.markdown(message["content"])
# If this is a user message and there are logs associated with it
# Display logs AFTER the user message but BEFORE the next assistant message
if message["role"] == "user" and "log_index" in message and i < len(st.session_state.messages) - 1:
log_index = message["log_index"]
if log_index < len(st.session_state.workflow_logs):
with st.expander("View Workflow Execution Logs", expanded=False):
st.code(st.session_state.workflow_logs[log_index], language="text")
# Accept user input
if prompt := st.chat_input("Ask a question about your documents..."):
# Add user message to chat history with placeholder for log index
log_index = len(st.session_state.workflow_logs)
st.session_state.messages.append({"role": "user", "content": prompt, "log_index": log_index})
# Display user message in chat message container
with st.chat_message("user"):
st.markdown(prompt)
if st.session_state.workflow:
# Run the async workflow
result = asyncio.run(run_workflow(prompt))
# Display the workflow logs in an expandable section OUTSIDE and BEFORE the assistant chat bubble
if log_index < len(st.session_state.workflow_logs):
with st.expander("View Workflow Execution Logs", expanded=False):
st.code(st.session_state.workflow_logs[log_index], language="text")
# Display assistant response in chat message container
with st.chat_message("assistant"):
if st.session_state.workflow:
message_placeholder = st.empty()
full_response = ""
result = result.response
# Stream the response word by word
words = result.split()
for i, word in enumerate(words):
full_response += word + " "
message_placeholder.markdown(full_response + "▌")
# Add a delay between words
if i < len(words) - 1: # Don't delay after the last word
time.sleep(0.1)
# Display final response without cursor
message_placeholder.markdown(full_response)
else:
full_response = "Please upload a document first to initialize the workflow."
st.markdown(full_response)
# Add assistant response to chat history
st.session_state.messages.append({"role": "assistant", "content": full_response})
corrective-rag/workflow.py
+169
View File
@@ -0,0 +1,169 @@
import os
from typing import Optional, Any
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.core.llms import LLM
from llama_index.llms.openai import OpenAI
from llama_index.core.base.base_retriever import BaseRetriever
from llama_index.tools.linkup_research.base import LinkupToolSpec
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,
linkup_api_key: str,
llm: Optional[LLM] = None,
**kwargs: Any
) -> None:
"""Init params."""
super().__init__(**kwargs)
self.index = index
self.linkup_tool = LinkupToolSpec(
api_key=linkup_api_key,
depth="deep", # or "deep"
output_type="searchResults", # or "sourcedAnswer" or "structured"
)
self.llm = 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", {})
if query_str is None:
return None
retriever: BaseRetriever = self.index.as_retriever(**retriever_kwargs)
result = retriever.retrieve(query_str)
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")
relevancy_results = []
for node in retrieved_nodes:
prompt = DEFAULT_RELEVANCY_PROMPT_TEMPLATE.format(context_str=node.text, query_str=query_str)
relevancy = self.llm.complete(prompt)
relevancy_results.append(relevancy.text.lower().strip())
relevant_texts = [
retrieved_nodes[i].text
for i, result in enumerate(relevancy_results)
if result == "yes"
]
relevant_text = "\n".join(relevant_texts)
if "no" in relevancy_results:
return WebSearchEvent(relevant_text=relevant_text)
else:
return QueryEvent(relevant_text=relevant_text, search_text="")
@step
async def web_search(
self, ctx: Context, ev: WebSearchEvent
) -> QueryEvent:
"""Search the transformed query with Tavily API."""
# 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)
result = self.llm.complete(prompt)
transformed_query_str = result.text
# Conduct a search with the transformed query string and collect the results.
search_results = self.linkup_tool.search(transformed_query_str).results
search_text = "\n".join([result.content for result in search_results])
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")
documents = [Document(text=relevant_text + "\n" + search_text)]
index = SummaryIndex.from_documents(documents)
query_engine = index.as_query_engine()
result = query_engine.query(query_str)
return StopEvent(result=result)