134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
import os
|
|
from dotenv import load_dotenv
|
|
from langchain_ollama import OllamaLLM
|
|
from langchain.agents import initialize_agent, AgentType
|
|
from langchain_community.tools.tavily_search import TavilySearchResults
|
|
from langchain.tools import tool
|
|
from vectorstore import get_vectorstore
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
|
|
if not TAVILY_API_KEY:
|
|
raise ValueError("TAVILY_API_KEY not found in .env file")
|
|
|
|
# Configuration
|
|
PERSIST_DIRECTORY = "./chroma_db"
|
|
DOCUMENTS_DIR = "./documents"
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
LLM_MODEL = "llama3"
|
|
|
|
# Initialize or load vectorstore
|
|
vectorstore = get_vectorstore()
|
|
|
|
# Define tools
|
|
@tool
|
|
def search_local_kb(query: str) -> str:
|
|
"""Search local knowledge base using ChromaDB."""
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
|
|
docs = retriever.get_relevant_documents(query)
|
|
return "\n\n".join([doc.page_content for doc in docs])
|
|
|
|
@tool
|
|
def web_search(query: str) -> str:
|
|
"""Search the web using Tavily."""
|
|
search = TavilySearchResults(tavily_api_key=TAVILY_API_KEY, max_results=3)
|
|
results = search.run(query)
|
|
return "\n\n".join([result["content"] for result in results])
|
|
|
|
tools = [search_local_kb, web_search]
|
|
|
|
# Initialize LLM
|
|
llm = OllamaLLM(model=LLM_MODEL)
|
|
|
|
# Create agent with routing logic
|
|
agent = initialize_agent(
|
|
tools=tools,
|
|
llm=llm,
|
|
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
|
verbose=True,
|
|
handle_parsing_errors=True,
|
|
system_message="""You are an AI agent that must choose between two tools:
|
|
1. search_local_kb: For questions about local documents (concepts, notes, stored information)
|
|
2. web_search: For questions requiring up-to-date information from the internet
|
|
|
|
When answering, ALWAYS specify your information source:
|
|
- If using local knowledge base: [Source: chromadb]
|
|
- If using web search: [Source: tavily]
|
|
|
|
Be concise and accurate in your responses."""
|
|
)
|
|
|
|
# Chat interface
|
|
def main():
|
|
print("RAG Agent with ChromaDB and Tavily Web Search")
|
|
print("Type 'exit' to quit\n")
|
|
|
|
while True:
|
|
query = input("Запрос: ").strip()
|
|
if query.lower() == "exit":
|
|
break
|
|
|
|
if not query:
|
|
continue
|
|
|
|
try:
|
|
response = agent.run(query)
|
|
print(f"\n{response}\n")
|
|
except Exception as e:
|
|
print(f"Error: {str(e)}\n")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
```
|
|
|
|
```python
|
|
# main.py
|
|
|
|
import os
|
|
from dotenv import load_dotenv
|
|
from langchain_community.tools.tavily_search import TavilySearchResults
|
|
from langchain.tools import tool
|
|
from vectorstore import get_vectorstore
|
|
|
|
load_dotenv()
|
|
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
|
|
if not TAVILY_API_KEY:
|
|
raise ValueError("TAVILY_API_KEY not found in .env file")
|
|
|
|
vectorstore = get_vectorstore()
|
|
|
|
@tool
|
|
def search_local_kb(query: str) -> str:
|
|
"""Search local knowledge base using ChromaDB."""
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
|
|
docs = retriever.get_relevant_documents(query)
|
|
return "\n\n".join([doc.page_content for doc in docs])
|
|
|
|
@tool
|
|
def web_search(query: str) -> str:
|
|
"""Search the web using Tavily."""
|
|
search = TavilySearchResults(tavily_api_key=TAVILY_API_KEY, max_results=3)
|
|
results = search.run(query)
|
|
return "\n\n".join([result["content"] for result in results])
|
|
```
|
|
|
|
```python
|
|
# tools.py
|
|
|
|
import os
|
|
from dotenv import load_dotenv
|
|
from langchain_community.tools.tavily_search import TavilySearchResults
|
|
from langchain.tools import tool
|
|
|
|
load_dotenv()
|
|
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
|
|
if not TAVILY_API_KEY:
|
|
raise ValueError("TAVILY_API_KEY not found in .env file")
|
|
|
|
@tool
|
|
def web_search(query: str) -> str:
|
|
"""Search the web using Tavily."""
|
|
search = TavilySearchResults(tavily_api_key=TAVILY_API_KEY, max_results=3)
|
|
results = search.run(query)
|
|
return "\n\n".join([result["content"] for result in results]) |