国产算力筑基，API为刃：大模型应用的积木式创新时代已至

def chunks_from_pdf(pdf_directory): """ Chunks all pdfs from a directory :param pdf_directory: directory of pdfs :return: list of chunks """ # fetching all pdfs from the directory and storing them as strings in a list docs = [] for file in glob.glob(pdf_directory + "/*.pdf"): loader = PyPDFLoader(file) doc = loader.load() docs.extend(doc) # split texts into chunks with overlap splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=500, chunk_overlap=100) splits = splitter.split_documents(docs) return splits def chunks_from_text(text_directory): """ Chunks all text files from a directory :param text_directory: directory of text files :return: list of chunks """ # fetch all txt files from the firectory and store them in a list loader = DirectoryLoader(text_directory, loader_cls=TextLoader) # , glob="**/*.txt") docs = loader.load() # split texts into chunks with overlap splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=500, chunk_overlap=100) splits = splitter.split_documents(docs) return splits def chunking(data_directory): """ Automatically calls the correct chunking function, either for pdfs or for txt files :param data_directory: directory of data, either ../pdf or ../text :return: result from the corresponding chunking function """ if data_directory == "../pdf": return chunks_from_pdf(data_directory) else: return chunks_from_text(data_directory)

from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.document_loaders import DirectoryLoader from langchain_community.document_loaders import PyPDFLoader from langchain_community.document_loaders import TextLoader from langchain_community.vectorstores import Chroma from langchain_core.embeddings import Embeddings class EmbeddingAgent(Embeddings): def __init__(self, model_name: str = "bge-m3"): self.api_url = API_URL self.model_name = model_name def embed_documents(self, texts: List[str]) -> List[List[float]]: non_empty_texts = [text for text in texts if text.strip()] if not non_empty_texts: return [[] for _ in texts] batch_size = 3 embeddings = [] for i in range(0, len(non_empty_texts), batch_size): batch = texts[i:i+batch_size] response = client.embeddings.create( input=batch[i], model=self.model_name, ) data = response.data for i in range(len(data)): embeddings.append(data[i].embedding) return embeddings def embed_query(self, text: str) -> List[float]: return self.embed_documents([text])[0]

def create_vector_store(db_directory, chunks, embedding): if not chunks: return None # create vector store and index vectorstore = Chroma.from_documents( documents=chunks, collection_name="rag_collection", embedding=embedding, persist_directory=db_directory ) return vectorstore.as_retriever() def fetch_vector_store(db_directory, embedding): print("Fetching vector store") vectorstore = Chroma( collection_name="chromemwah", embedding_function=embedding, persist_directory=db_directory ) return vectorstore.as_retriever() def retrieve(retrieving, question): documents = retrieving.get_relevant_documents(question) return documents

class ChatAgent: def __init__(self, model_name: str): self.api_url = API_URL self.model_name = model_name def invoke(self, messages: List[Dict[str, str]]) -> str: try: extra_headers = {} completion = client.chat.completions.create( extra_headers=extra_headers, model=self.model_name, messages=messages, stream=True, temperature=0.6, ) for chunk in completion: if len(chunk.choices) != 0: delta = chunk.choices[0].delta # If is thinking content, print it in gray if hasattr(delta, 'reasoning_content') and delta.reasoning_content: fullResponse += delta.reasoning_content elif delta.content: fullResponse += delta.content # Format response to match expected output result = [{ "message": { "content": fullResponse } }] except Exception as e: logger.error(f"Error with Gitee API: {e}") raise return result[0]['message']['content'] def generate(question, context, llm_model): system_prompt = ( "You are an assistant for question-answering tasks. Use the " "following pieces of retrieved context to answer the question. If " "you don't know the answer, just say that you don't know. Keep the " "answer concise, truthful, and informative. If you decide to use a " "source, you must mention in which document you found specific " "information. Sources are indicated in the context by " "[doc<doc_number>]." ) user_prompt = f"Question: {question} \nContext: {context}" messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ] return ChatAgent.invoke(messages)