201 lines
9.0 KiB
Python
201 lines
9.0 KiB
Python
# --- Dependencies ---
|
|
# pip install langchain langchain-core langchain-ollama faiss-cpu sentence-transformers
|
|
# pip install chromadb # Example if switching later
|
|
|
|
import datetime
|
|
import os
|
|
import argparse # Import argparse for command-line arguments
|
|
# (REMOVE LATER - LLM Related)
|
|
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
|
from langchain_community.vectorstores import FAISS
|
|
# (REMOVE LATER - LLM Related)
|
|
from langchain.prompts import PromptTemplate
|
|
# (REMOVE LATER - LLM Related)
|
|
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
|
|
# (REMOVE LATER - LLM Related)
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from langchain.schema import Document
|
|
|
|
# --- Config ---
|
|
FAISS_INDEX_PATH = "my_chatbot_memory_index_user_specific" # Directory for FAISS index
|
|
|
|
# --- Ollama LLM & Embeddings Setup ---
|
|
# (REMOVE LATER - LLM Related)
|
|
OLLAMA_LLM_MODEL = 'gemma3'
|
|
OLLAMA_EMBED_MODEL = 'nomic-embed-text' # Keep embeddings even if LLM removed
|
|
|
|
# Global variables for embeddings and vector store (initialized later)
|
|
embeddings = None
|
|
vectorstore = None
|
|
retriever = None
|
|
|
|
def initialize_memory():
|
|
"""Initializes the Ollama embeddings and loads/creates the FAISS vector store."""
|
|
global embeddings, vectorstore, retriever
|
|
|
|
if embeddings and vectorstore: # Prevent re-initialization
|
|
# print("Memory already initialized.") # Optional: uncomment for verbose logging
|
|
return
|
|
|
|
try:
|
|
embeddings = OllamaEmbeddings(model=OLLAMA_EMBED_MODEL)
|
|
print(f"Successfully initialized Ollama Embeddings: '{OLLAMA_EMBED_MODEL}'")
|
|
except Exception as e:
|
|
print(f"Error initializing Ollama Embeddings: {e}")
|
|
print(f"Ensure Ollama Embedding model is running & pulled ('ollama pull {OLLAMA_EMBED_MODEL}').")
|
|
raise e # Re-raise the exception so the caller knows initialization failed
|
|
|
|
try:
|
|
# Check if the actual index file exists, not just the directory
|
|
if os.path.exists(os.path.join(FAISS_INDEX_PATH, "index.faiss")):
|
|
print(f"Loading existing FAISS index from: {FAISS_INDEX_PATH}")
|
|
vectorstore = FAISS.load_local(
|
|
FAISS_INDEX_PATH,
|
|
embeddings,
|
|
allow_dangerous_deserialization=True
|
|
)
|
|
print("FAISS vector store loaded successfully.")
|
|
else:
|
|
print(f"No FAISS index found at {FAISS_INDEX_PATH}. Initializing new store.")
|
|
# FAISS needs at least one text to initialize.
|
|
vectorstore = FAISS.from_texts(
|
|
["system_placeholder_init"], # Use a placeholder text
|
|
embeddings,
|
|
metadatas=[{"user_id": "system_init"}] # Placeholder metadata
|
|
)
|
|
vectorstore.save_local(FAISS_INDEX_PATH)
|
|
print("New FAISS vector store initialized and saved.")
|
|
|
|
# Initialize retriever after vectorstore is loaded/created
|
|
retriever = vectorstore.as_retriever(search_kwargs=dict(k=5)) # Retrieve more initially
|
|
|
|
except Exception as e:
|
|
print(f"Error initializing/loading FAISS: {e}")
|
|
print("Check permissions or delete the index directory if corrupted.")
|
|
raise e # Re-raise the exception
|
|
|
|
def add_memory(user_id: str, user_input: str, assistant_output: str):
|
|
"""Adds a user query and assistant response to the memory for a specific user."""
|
|
global vectorstore
|
|
|
|
if not vectorstore:
|
|
print("Error: Memory not initialized. Call initialize_memory() first.")
|
|
return False
|
|
|
|
if user_input and assistant_output and user_id:
|
|
timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
docs_to_add = [
|
|
Document(page_content=f"{user_input}", metadata={"user_id": user_id, "role": "user", "timestamp": timestamp}),
|
|
Document(page_content=f"{assistant_output}", metadata={"user_id": user_id, "role": "assistant", "timestamp": timestamp})
|
|
]
|
|
try:
|
|
vectorstore.add_documents(docs_to_add)
|
|
vectorstore.save_local(FAISS_INDEX_PATH) # Persist index after adding
|
|
# print(f"Memory added for user {user_id}.") # Optional: uncomment for verbose logging
|
|
return True
|
|
except Exception as e:
|
|
print(f"Error adding document to FAISS for user {user_id}: {e}")
|
|
return False
|
|
else:
|
|
print("Warning: Attempted to add memory with missing user_id, user_input, or assistant_output.")
|
|
return False
|
|
|
|
def get_formatted_context(user_id: str, query: str, k: int = 3) -> str:
|
|
"""Retrieves relevant memories for a user based on a query and formats them."""
|
|
global retriever
|
|
|
|
if not retriever:
|
|
print("Error: Memory not initialized. Call initialize_memory() first.")
|
|
return "خطا: حافظه مقداردهی اولیه نشده است."
|
|
|
|
try:
|
|
current_retriever = vectorstore.as_retriever(search_kwargs=dict(k=k*2))
|
|
docs = current_retriever.invoke(query)
|
|
|
|
user_specific_docs = [doc for doc in docs if doc.metadata.get("user_id") == user_id]
|
|
|
|
formatted = []
|
|
for doc in user_specific_docs[:k]:
|
|
content = doc.page_content
|
|
role = doc.metadata.get("role", "unknown")
|
|
if content and content != "system_placeholder_init":
|
|
formatted.append(f"{role.capitalize()}: {content}")
|
|
|
|
return "\\n".join(formatted) if formatted else "هیچ خاطره مرتبطی از جلسات گذشته برای این کاربر یافت نشد."
|
|
|
|
except Exception as e:
|
|
print(f"Error retrieving context from FAISS for user {user_id}: {e}")
|
|
return "خطا در بازیابی زمینه حافظه."
|
|
|
|
# --- Main Execution (Command-Line Interface with Logging) ---
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Memory module CLI tool with logging")
|
|
parser.add_argument("--action", required=True, choices=['add', 'get'], help="Action to perform: add or get memory.")
|
|
parser.add_argument("--user-id", default="default_user", help="User ID for the memory operation (optional, defaults to 'default_user')")
|
|
parser.add_argument("--input", required=True, help="Input text (user query for 'get', user message for 'add')")
|
|
parser.add_argument("--output", help="Output text (assistant response, required for 'add' action)")
|
|
parser.add_argument("--k", type=int, default=3, help="Number of memories to retrieve for 'get' action (optional, defaults to 3)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
log_filename = f"cli_log_{args.user_id}_{args.action}.txt"
|
|
log_content = []
|
|
log_content.append(f"Timestamp: {datetime.datetime.now()}")
|
|
log_content.append(f"Action: {args.action}")
|
|
log_content.append(f"User ID: {args.user_id}")
|
|
log_content.append(f"Input: {args.input}")
|
|
if args.action == 'add':
|
|
log_content.append(f"Output (for add): {args.output}")
|
|
if args.action == 'get':
|
|
log_content.append(f"K (for get): {args.k}")
|
|
log_content.append("---")
|
|
|
|
try:
|
|
# Initialize memory system first
|
|
print("Initializing memory system...")
|
|
initialize_memory()
|
|
print("Memory system initialized.")
|
|
|
|
# Perform requested action
|
|
if args.action == 'add':
|
|
if not args.output:
|
|
error_msg = "Error: --output is required for action 'add'."
|
|
print(error_msg)
|
|
log_content.append(f"Result: FAILURE - {error_msg}")
|
|
else:
|
|
if add_memory(args.user_id, args.input, args.output):
|
|
result_msg = f"Successfully added memory for user '{args.user_id}'."
|
|
print(result_msg)
|
|
log_content.append(f"Result: SUCCESS - {result_msg}")
|
|
else:
|
|
result_msg = f"Failed to add memory for user '{args.user_id}'. Check logs for details."
|
|
print(result_msg)
|
|
log_content.append(f"Result: FAILURE - {result_msg}")
|
|
|
|
elif args.action == 'get':
|
|
context = get_formatted_context(args.user_id, args.input, args.k)
|
|
# Print the raw context directly to standard output for immediate feedback
|
|
print("--- Retrieved Context --- ")
|
|
print(context)
|
|
print("-------------------------")
|
|
log_content.append("Result:")
|
|
log_content.append(context)
|
|
|
|
except Exception as e:
|
|
error_msg = f"An error occurred during execution: {e}"
|
|
print(error_msg)
|
|
log_content.append(f"Result: FAILURE - {error_msg}")
|
|
import traceback
|
|
log_content.append(traceback.format_exc())
|
|
|
|
finally:
|
|
# Write the log file regardless of success or failure (if possible)
|
|
try:
|
|
with open(log_filename, 'w', encoding='utf-8') as f:
|
|
f.write("\\n".join(log_content))
|
|
print(f"--- Log saved to: {log_filename} ---")
|
|
except Exception as log_e:
|
|
print(f"Error writing log file {log_filename}: {log_e}")
|
|
|
|
# --- End of Script ---
|