1# You may need to add your working directory to the Python path. To do so, uncomment the following lines of code
2# import sys
3# sys.path.append("/Path/to/directory/agentic-framework") # Replace with your directory path
4
5import logging
6import uuid
7
8from chromadb import EphemeralClient
9from langchain_community.embeddings import OpenAIEmbeddings
10from langchain_community.vectorstores import Chroma
11from langchain_text_splitters import RecursiveCharacterTextSplitter
12
13from baf import nlp
14from baf.core.agent import Agent
15from baf.core.session import Session
16from baf.exceptions.logger import logger
17from baf.library.transition.events.base_events import ReceiveFileEvent
18from baf.nlp.llm.llm_huggingface_api import LLMHuggingFaceAPI
19from baf.nlp.llm.llm_openai_api import LLMOpenAI
20from baf.nlp.llm.llm_replicate_api import LLMReplicate
21from baf.nlp.rag.rag import RAGMessage, RAG
22from baf.library.transition.events.base_events import ReceiveTextEvent, ReceiveMessageEvent, ReceiveFileEvent
23
24# Configure the logging module (optional)
25logger.setLevel(logging.INFO)
26
27# Create the agent
28agent = Agent('rag_agent')
29# Load agent properties stored in a dedicated file
30agent.load_properties('config.yaml')
31# Define the platform your agent will use
32websocket_platform = agent.use_websocket_platform(use_ui=True)
33
34#To keep RAG as session scoped or not
35SESSION_SCOPED=True
36
37
38# Create text splitter (RAG creates a vector for each chunk)
39splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
40# Create the LLM (for the answer generation)
41gpt = LLMOpenAI(
42 agent=agent,
43 name='gpt-4o-mini',
44 parameters={},
45 num_previous_messages=10
46)
47
48# Other example LLM
49
50# gemma = LLMHuggingFace(agent=agent, name='google/gemma-2b-it', parameters={'max_new_tokens': 1}, num_previous_messages=10)
51# llama = LLMHuggingFaceAPI(agent=agent, name='meta-llama/Meta-Llama-3.1-8B-Instruct', parameters={}, num_previous_messages=10)
52# mixtral = LLMReplicate(agent=agent, name='mistralai/mixtral-8x7b-instruct-v0.1', parameters={}, num_previous_messages=10)
53
54# Create the RAG
55rag = RAG(
56 agent=agent,
57 vector_store=vector_store,
58 splitter=splitter,
59 llm_name='gpt-4o-mini',
60 k=4,
61 num_previous_messages=0
62)
63# Uncomment the relevant line to directly provide text to RAG instead of uploading at runtime
64# rag.load_documents('C:/Users/chidambaram/Downloads/example', formats=['docx']) # read all the files of the given format in the directory
65# rag.add_text('raw text to index and base your RAG on.') # add a text string directly
66
67# If not uploading at runtime, comment the following function, L84, L87 and L88, and uncomment L85.
68def upload_file_body(session: Session):
69 rag.add_file(session.event.file)
70
71
72# STATES
73
74initial_state = agent.new_state('initial_state', initial=True)
75upload_file_state = agent.new_state('upload_file_state')
76ask_question_state = agent.new_state('ask_question_state')
77
78
79# STATES BODIES' DEFINITION + TRANSITIONS
80
81def initial_body(session: Session):
82 session.reply('Hi!, upload a file to use as RAG context. If already done, ask your question.')
83
84
85initial_state.set_body(initial_body)
86initial_state.when_event(ReceiveFileEvent()).go_to(upload_file_state)
87# initial_state.when_no_intent_matched().go_to(ask_question_state)
88
89upload_file_state.set_body(upload_file_body)
90upload_file_state.when_no_intent_matched().go_to(ask_question_state)
91
92
93def ask_question_body(session: Session):
94 rag_message: RAGMessage = session.run_rag(session.event.message)
95 websocket_platform.reply_rag(session, rag_message)
96
97
98ask_question_state.set_body(ask_question_body)
99ask_question_state.go_to(initial_state)
100
101
102# RUN APPLICATION
103
104if __name__ == '__main__':
105 agent.run()