Hacker and Geeker's Way

from llama_index.llms.openai_like import OpenAILike
from llama_index.core.base.llms.types import ChatMessage, MessageRole
from llama_index.core import PromptTemplate

llm = OpenAILike(
    model="llama3",
    api_base="you-local-llama3-api",
    api_key="fake_key",
    is_chat_model=True,
)
prompt_str = "Please generate related movies to {movie_name}"
prompt_tmpl = PromptTemplate(prompt_str)
response = llm.chat(
    [
        ChatMessage(
            role=MessageRole.SYSTEM,
            content="You are a helpful assistant.",
        ),
        ChatMessage(
            role=MessageRole.USER,
            content=prompt_tmpl.format(movie_name="Avengers"),
        ),
    ]
)
print(f"response: {response}")

# 显示结果
response: assistant: Here are some movie recommendations that are similar to the Avengers franchise:

1. **Guardians of the Galaxy** (2014) - Another Marvel superhero team-up film, with a fun and quirky tone.
2. **The Justice League** (2017) - A DC Comics adaptation featuring iconic superheroes like Superman, Batman, Wonder Woman, and more.
......

from llama_index.vector_stores.elasticsearch import ElasticsearchStore

es = ElasticsearchStore(
    index_name="my_index",
    es_url="http://localhost:9200",
)

from llama_index.vector_stores.elasticsearch import ElasticsearchStore
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.core.node_parser import SentenceSplitter
from llms import CustomEmbeddings

store = ElasticsearchStore(
    index_name="avengers",
    es_url="http://localhost:9200",
)
documents = SimpleDirectoryReader("./data").load_data()
node_parser = SentenceSplitter(chunk_size=256, chunk_overlap=50)
storage_context = StorageContext.from_defaults(vector_store=store)
embed_model = CustomEmbeddings(
    model="BAAI/bge-base-en-v1.5", url="http://localhost:6006"
)
VectorStoreIndex.from_documents(
    documents,
    transformations=[node_parser],
    embed_model=embed_model,
    storage_context=storage_context,
)

from llama_index.vector_stores.elasticsearch import AsyncBM25Strategy
from llama_index.core import Settings

text_store = ElasticsearchStore(
    index_name="avengers",
    es_url="http://localhost:9200",
    retrieval_strategy=AsyncBM25Strategy(),
)
Settings.embed_model = embed_model
text_index = VectorStoreIndex.from_vector_store(
    vector_store=text_store,
)
text_retriever = text_index.as_retriever(similarity_top_k=2)

from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy, AsyncSparseVectorStrategy

vector_store = ElasticsearchStore(
    index_name="avengers",
    es_url="http://localhost:9200",
    retrieval_strategy=AsyncDenseVectorStrategy(),
    # retrieval_strategy=AsyncSparseVectorStrategy(model_id=".elser_model_2"),
)
Settings.embed_model = embed_model
vector_index = VectorStoreIndex.from_vector_store(
    vector_store=vector_store,
)
vector_retriever = vector_index.as_retriever(similarity_top_k=2)

from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy

vector_store = ElasticsearchStore(
    index_name="avengers",
    es_url="http://localhost:9200",
    retrieval_strategy=AsyncDenseVectorStrategy(hybrid=True),
)

from typing import List
from llama_index.core.schema import NodeWithScore

def fuse_results(results_dict, similarity_top_k: int = 2):
    """Fuse results."""
    k = 60.0
    fused_scores = {}
    text_to_node = {}

    # 计算倒数排名分数
    for nodes_with_scores in results_dict.values():
        for rank, node_with_score in enumerate(
            sorted(
                nodes_with_scores, key=lambda x: x.score or 0.0, reverse=True
            )
        ):
            text = node_with_score.node.get_content()
            text_to_node[text] = node_with_score
            if text not in fused_scores:
                fused_scores[text] = 0.0
            fused_scores[text] += 1.0 / (rank + k)

    # 结果按分数排序
    reranked_results = dict(
        sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    )

    # 结果还原为节点集合
    reranked_nodes: List[NodeWithScore] = []
    for text, score in reranked_results.items():
        reranked_nodes.append(text_to_node[text])
        reranked_nodes[-1].score = score

    return reranked_nodes[:similarity_top_k]

from tqdm.asyncio import tqdm

def run_queries(query, retrievers):
    """Run query against retrievers."""
    tasks = []
    for i, retriever in enumerate(retrievers):
        tasks.append(retriever.aretrieve(query))

    task_results = await tqdm.gather(*tasks)

    results_dict = {}
    for i, query_result in enumerate(task_results):
        results_dict[(query, i)] = query_result

    return results_dict

from typing import List
from llama_index.core import QueryBundle
from llama_index.core.retrievers import BaseRetriever
from llama_index.core.schema import NodeWithScore
import asyncio

class FusionRetriever(BaseRetriever):
    """Ensemble retriever with fusion."""

    def __init__(
        self,
        retrievers: List[BaseRetriever],
        similarity_top_k: int = 2,
    ) -> None:
        """Init params."""
        self._retrievers = retrievers
        self._similarity_top_k = similarity_top_k
        super().__init__()

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Retrieve."""
        results = asyncio.run(
            run_queries(query_bundle.query_str, self._retrievers)
        )
        final_results = fuse_results(results, similarity_top_k=self._similarity_top_k)
        return final_results

fusion_retriever = FusionRetriever(
    [text_retriever, vector_retriever], similarity_top_k=2
)
question = "Which two members of the Avengers created Ultron?"
nodes = fusion_retriever.retrieve(question)
for node in nodes:
    print("-" * 50)
    print(f"node content: {node.text[:100]}...")
    print(f"node score: {node.score}\n")

# 显示结果
--------------------------------------------------
node content: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...
node score: 0.03306010928961749

--------------------------------------------------
node content: Thor departs to consult with Dr. Erik Selvig on the apocalyptic future he saw in his hallucination, ...
node score: 0.016666666666666666

from llama_index.core.query_engine import RetrieverQueryEngine

rerank = CustomRerank(
    model="BAAI/bge-reranker-base", url="http://localhost:7007", top_n=2
)
Settings.llm = llm
query_engine = RetrieverQueryEngine(fusion_retriever, node_postprocessors=[rerank])
response = query_engine.query(question)
print(f"response: {response}")
for node in response.source_nodes:
    print("-" * 50)
    print(f"node content: {node.text[:100]}...")
    print(f"node score: {node.score}\n")

# 显示结果
response: Tony Stark and Bruce Banner.
--------------------------------------------------
node content: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...
node score: 0.8329173

--------------------------------------------------
node content: Thor departs to consult with Dr. Erik Selvig on the apocalyptic future he saw in his hallucination, ...
node score: 0.24689633