from sentence_transformers import SentenceTransformer

def get_embedding_function():
    # Load the local embedding model
    model = SentenceTransformer('all-MiniLM-L6-v2')  # You can choose another model from Hugging Face

    # Create an embedding function with `embed_documents` and `embed_query`
    class EmbeddingsWrapper:
        def embed_documents(self, texts):
            """Embed a list of documents (texts)."""
            embeddings = model.encode(texts, convert_to_tensor=False)
            # Convert to list to avoid ambiguity with array truth values
            return [embedding.tolist() if hasattr(embedding, "tolist") else embedding for embedding in embeddings]

        def embed_query(self, text):
            """Embed a single query."""
            embedding = model.encode([text], convert_to_tensor=False)[0]
            return embedding.tolist() if hasattr(embedding, "tolist") else embedding

    return EmbeddingsWrapper()