Spaces:
Running
Running
File size: 4,838 Bytes
2d671a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
"""RAG pipeline for Graphics Guide / Design Assistant"""
from typing import Iterator, Optional, List, Tuple
from .vectorstore import GraphicsVectorStore, create_vectorstore
from .llm_client import InferenceProviderClient, create_llm_client
from .prompts import (
SYSTEM_PROMPT,
DESIGN_PROMPT,
get_design_prompt
)
class GraphicsDesignPipeline:
"""RAG pipeline for generating graphics and design recommendations"""
def __init__(
self,
vectorstore: Optional[GraphicsVectorStore] = None,
llm_client: Optional[InferenceProviderClient] = None,
retrieval_k: int = 5
):
"""
Initialize the RAG pipeline
Args:
vectorstore: Vector store instance (creates default if None)
llm_client: LLM client instance (creates default if None)
retrieval_k: Number of document chunks to retrieve for context
"""
self.vectorstore = vectorstore or create_vectorstore()
self.llm_client = llm_client or create_llm_client()
self.retrieval_k = retrieval_k
def retrieve_documents(self, query: str, k: Optional[int] = None) -> List:
"""
Retrieve relevant document chunks for a query
Args:
query: User's design query
k: Number of documents to retrieve (uses default if None)
Returns:
List of relevant document chunks
"""
k = k or self.retrieval_k
return self.vectorstore.similarity_search(query, k=k)
def generate_recommendations(
self,
query: str,
stream: bool = False
) -> str | Iterator[str]:
"""
Generate design recommendations for a query
Args:
query: User's design query
stream: Whether to stream the response
Returns:
Generated recommendations (string or iterator)
"""
# Retrieve relevant documents
relevant_docs = self.retrieve_documents(query)
# Format documents for context
context = self.vectorstore.format_documents_for_context(relevant_docs)
# Generate prompt
prompt_template = get_design_prompt()
full_prompt = prompt_template.format(query=query, context=context)
# Generate response
if stream:
return self.llm_client.generate_stream(
prompt=full_prompt,
system_prompt=SYSTEM_PROMPT
)
else:
return self.llm_client.generate(
prompt=full_prompt,
system_prompt=SYSTEM_PROMPT
)
def chat(
self,
message: str,
history: Optional[List[Tuple[str, str]]] = None,
stream: bool = False
) -> str | Iterator[str]:
"""
Handle a chat message with conversation history
Args:
message: User's message
history: Conversation history as list of (user_msg, assistant_msg) tuples
stream: Whether to stream the response
Returns:
Generated response (string or iterator)
"""
# For now, treat each message as a new design query
# In the future, could implement follow-up handling
return self.generate_recommendations(message, stream=stream)
def get_relevant_examples(
self,
query: str,
k: int = 5
) -> List[dict]:
"""
Get relevant examples and knowledge with metadata
Args:
query: Design query
k: Number of examples to recommend
Returns:
List of document dictionaries with metadata
"""
docs = self.retrieve_documents(query, k=k)
examples = []
for doc in docs:
example = {
"source": doc.metadata.get("source_id", "Unknown"),
"source_type": doc.metadata.get("source_type", "N/A"),
"page": doc.metadata.get("page_number"),
"content": doc.page_content,
"similarity": doc.metadata.get("similarity")
}
examples.append(example)
return examples
def create_pipeline(
retrieval_k: int = 5,
model: str = "meta-llama/Llama-3.1-8B-Instruct",
temperature: float = 0.2
) -> GraphicsDesignPipeline:
"""
Factory function to create a configured RAG pipeline
Args:
retrieval_k: Number of documents to retrieve
model: LLM model identifier
temperature: LLM temperature
Returns:
Configured GraphicsDesignPipeline
"""
vectorstore = create_vectorstore()
llm_client = create_llm_client(model=model, temperature=temperature)
return GraphicsDesignPipeline(
vectorstore=vectorstore,
llm_client=llm_client,
retrieval_k=retrieval_k
)
|