from feast import FeatureStore

# From repo path (looks for feature_store.yaml)
store = FeatureStore(repo_path="path/to/feature_repo")

# From config object
from feast.repo_config import RepoConfig
store = FeatureStore(config=RepoConfig(
    project="my_project",
    registry="data/registry.db",
    provider="local",
    online_store={"type": "sqlite", "path": "data/online.db"},
))

# From explicit YAML path
from pathlib import Path
store = FeatureStore(fs_yaml_file=Path("custom/feature_store.yaml"))

result = store.get_online_features(
    features=[
        "driver_hourly_stats:conv_rate",
        "driver_hourly_stats:acc_rate",
        "driver_hourly_stats:avg_daily_trips",
    ],
    entity_rows=[
        {"driver_id": 1001},
        {"driver_id": 1002},
    ],
)

feature_dict = result.to_dict()
feature_df = result.to_df()

result = store.get_online_features(
    features=driver_ranking_service,
    entity_rows=[{"driver_id": 1001}],
)

import pandas as pd
from datetime import datetime

entity_df = pd.DataFrame({
    "driver_id": [1001, 1002, 1003],
    "event_timestamp": [
        datetime(2023, 6, 1),
        datetime(2023, 6, 15),
        datetime(2023, 7, 1),
    ],
})

training_df = store.get_historical_features(
    entity_df=entity_df,
    features=[
        "driver_hourly_stats:conv_rate",
        "driver_hourly_stats:acc_rate",
    ],
).to_df()

training_df = store.get_historical_features(
    entity_df=entity_df,
    features=driver_ranking_service,
).to_df()

store.push(
    push_source_name="driver_push",
    df=pd.DataFrame({
        "driver_id": [1001],
        "trips_today": [15],
        "event_timestamp": [datetime.utcnow()],
    }),
)

store.write_to_online_store(
    feature_view_name="driver_hourly_stats",
    df=features_df,
)

store.write_to_offline_store(
    feature_view_name="driver_hourly_stats",
    df=features_df,
)

from feast import Entity, FeatureView, Field, FileSource
from feast.types import Array, Float32, String

passage_entity = Entity(name="passage_id", join_keys=["passage_id"])

wiki_passages = FeatureView(
    name="wiki_passages",
    entities=[passage_entity],
    schema=[
        Field(name="passage_text", dtype=String),
        Field(
            name="embedding",
            dtype=Array(Float32),
            vector_index=True,
            vector_length=384,
            vector_search_metric="COSINE",
        ),
    ],
    source=passages_source,
    online=True,
)

# v1 API
results = store.retrieve_online_documents(
    feature="wiki_passages:embedding",
    query=query_embedding_vector,
    top_k=5,
)

# v2 API (supports text, vector, and image queries)
results = store.retrieve_online_documents_v2(
    feature_view_name="wiki_passages",
    query_string="What is machine learning?",
    top_k=5,
)

from feast.rag_retriever import FeastRAGRetriever
from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM

question_tokenizer = AutoTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
question_encoder = AutoModel.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
generator_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
generator_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large")

retriever = FeastRAGRetriever(
    question_encoder_tokenizer=question_tokenizer,
    question_encoder=question_encoder,
    generator_tokenizer=generator_tokenizer,
    generator_model=generator_model,
    feast_repo_path="path/to/feature_repo",
    feature_view="wiki_passages",
    features=["passage_text", "embedding"],
    search_type="vector",   # "text", "vector", or "hybrid"
    id_field="passage_id",
    text_field="passage_text",
)

doc_embeddings, doc_ids, doc_dicts = retriever.retrieve(
    question_input_ids=question_tokenizer("What is ML?", return_tensors="pt")["input_ids"],
    n_docs=5,
)

answer = retriever.generate_answer(
    query="What is machine learning?",
    top_k=5,
    max_new_tokens=200,
)
print(answer)

from feast.vector_store import FeastVectorStore

vector_store = FeastVectorStore(feast_repo_path="path/to/feature_repo")

results = vector_store.query(
    query_vector=embedding_list,
    top_k=10,
)