国产算力筑基，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)