Spaces:
Sleeping
Sleeping
File size: 5,770 Bytes
6466c00 |
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 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 |
"""RAG pipeline for OSINT investigation assistant"""
from typing import Iterator, Optional, List, Tuple
from .vectorstore import OSINTVectorStore, create_vectorstore
from .llm_client import InferenceProviderClient, create_llm_client
from .prompts import (
SYSTEM_PROMPT,
INVESTIGATION_PROMPT,
get_investigation_prompt
)
class OSINTInvestigationPipeline:
"""RAG pipeline for generating OSINT investigation methodologies"""
def __init__(
self,
vectorstore: Optional[OSINTVectorStore] = 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 tools 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_tools(self, query: str, k: Optional[int] = None) -> List:
"""
Retrieve relevant OSINT tools for a query
Args:
query: User's investigation query
k: Number of tools to retrieve (uses default if None)
Returns:
List of relevant tool documents
"""
k = k or self.retrieval_k
return self.vectorstore.similarity_search(query, k=k)
def generate_methodology(
self,
query: str,
stream: bool = False
) -> str | Iterator[str]:
"""
Generate investigation methodology for a query
Args:
query: User's investigation query
stream: Whether to stream the response
Returns:
Generated methodology (string or iterator)
"""
# Retrieve relevant tools
relevant_tools = self.retrieve_tools(query)
# Format tools for context
context = self.vectorstore.format_tools_for_context(relevant_tools)
# Generate prompt
prompt_template = get_investigation_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 investigation query
# In the future, could implement follow-up handling
return self.generate_methodology(message, stream=stream)
def get_tool_recommendations(
self,
query: str,
k: int = 5
) -> List[dict]:
"""
Get tool recommendations with metadata
Args:
query: Investigation query
k: Number of tools to recommend
Returns:
List of tool dictionaries with metadata
"""
docs = self.retrieve_tools(query, k=k)
tools = []
for doc in docs:
tool = {
"name": doc.metadata.get("name", "Unknown"),
"category": doc.metadata.get("category", "N/A"),
"cost": doc.metadata.get("cost", "N/A"),
"url": doc.metadata.get("url", "N/A"),
"description": doc.page_content,
"details": doc.metadata.get("details", "N/A")
}
tools.append(tool)
return tools
def search_tools_by_category(
self,
category: str,
k: int = 10
) -> List[dict]:
"""
Search tools by category
Args:
category: Tool category (e.g., "Archiving", "Social Media")
k: Number of tools to return
Returns:
List of tool dictionaries
"""
docs = self.vectorstore.similarity_search(
query=category,
k=k,
filter_category=category
)
tools = []
for doc in docs:
tool = {
"name": doc.metadata.get("name", "Unknown"),
"category": doc.metadata.get("category", "N/A"),
"cost": doc.metadata.get("cost", "N/A"),
"url": doc.metadata.get("url", "N/A"),
"description": doc.page_content
}
tools.append(tool)
return tools
def create_pipeline(
retrieval_k: int = 5,
model: str = "meta-llama/Llama-3.1-8B-Instruct",
temperature: float = 0.2
) -> OSINTInvestigationPipeline:
"""
Factory function to create a configured RAG pipeline
Args:
retrieval_k: Number of tools to retrieve
model: LLM model identifier
temperature: LLM temperature
Returns:
Configured OSINTInvestigationPipeline
"""
vectorstore = create_vectorstore()
llm_client = create_llm_client(model=model, temperature=temperature)
return OSINTInvestigationPipeline(
vectorstore=vectorstore,
llm_client=llm_client,
retrieval_k=retrieval_k
)
|