Chroma

import os
import pandas as pd
from openai import OpenAI

import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction

from honeyhive import HoneyHiveTracer, trace


HoneyHiveTracer.init(
    api_key="MY_HONEYHIVE_API_KEY", # paste your API key here
    project="MY_HONEYHIVE_PROJECT_NAME", # paste your project name here
)

client = OpenAI()
embedding_function = OpenAIEmbeddingFunction(api_key=os.getenv("OPENAI_API_KEY"))

claim_df = pd.read_json("scifact_claims.jsonl", lines=True)
corpus_df = pd.read_json("scifact_corpus.jsonl", lines=True)
corpus_df = corpus_df.sample(10) # comment this out to use full corpus

chroma_client = chromadb.Client()
scifact_corpus_collection = chroma_client.create_collection(
    name="scifact_corpus", embedding_function=embedding_function
)

batch_size = 100

for i in range(0, len(corpus_df), batch_size):
    batch_df = corpus_df[i: i + batch_size]
    scifact_corpus_collection.add(
        ids=batch_df["doc_id"]
        .apply(lambda x: str(x))
        .tolist(),  # Chroma takes string IDs.
        documents=(
            batch_df["title"] + ". " + batch_df["abstract"].apply(lambda x: " ".join(x))
        ).to_list(),  # We concatenate the title and abstract.
        metadatas=[
            {"structured": structured}
            for structured in batch_df["structured"].to_list()
        ],  # We also store the metadata, though we don't use it in this example.
    )


def build_prompt_with_context(claim, context):
    return [
        {
            "role": "system",
            "content": "I will ask you to assess whether a particular scientific claim, based on evidence provided. "
            + "Output only the text 'True' if the claim is true, 'False' if the claim is false, or 'NEE' if there's "
            + "not enough evidence.",
        },
        {
            "role": "user",
            "content": f""""
The evidence is the following:

{' '.join(context)}

Assess the following claim on the basis of the evidence. Output only the text 'True' if the claim is true,
'False' if the claim is false, or 'NEE' if there's not enough evidence. Do not output any other text.

Claim:
{claim}

Assessment:
""",
        },
    ]


@trace
def assess_claims(claims):
    claim_query_result = scifact_corpus_collection.query(
        query_texts=claims, include=["documents", "distances"], n_results=3
    )
    responses = []
    # Query the OpenAI API
    for claim, context in zip(claims, claim_query_result["documents"]):
        # If no evidence is provided, return NEE
        if len(context) == 0:
            responses.append("NEE")
            continue
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=build_prompt_with_context(claim=claim, context=context),
            max_tokens=3,
        )
        # Strip any punctuation or whitespace from the response
        formatted_response = response.choices[0].message.content.strip("., ")
        print("Claim: ", claim)
        print("Response: ", formatted_response)
        responses.append(formatted_response)

    return responses


samples = claim_df.sample(2)
assess_claims(samples["claim"].tolist())

import { ChromaClient, OpenAIEmbeddingFunction } from "chromadb";
import OpenAI from "openai";
import fs from "fs";
import { HoneyHiveTracer } from 'honeyhive';

const tracer = await HoneyHiveTracer.init({
    apiKey: 'MY_HONEYHIVE_API_KEY',
    project: 'MY_HONEYHIVE_PROJECT_NAME',
    sessionName: 'chromadb',
});

const openai_client = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
const client = new ChromaClient();

const embeddingFunction = new OpenAIEmbeddingFunction({
    openai_api_key: process.env.OPENAI_API_KEY ?? "",
});

const scifactCorpusCollection = client.getOrCreateCollection({
    name: "scifact_corpus",
    embeddingFunction,
});

interface SciFactData {
    doc_id: number;
    title: string;
    abstract: string[];
    structured: any;
    claim?: string;
}

const claimData: SciFactData[] = fs
    .readFileSync("scifact_claims.jsonl")
    .toString()
    .split("\n")
    .map((each) => {
        try {
            return JSON.parse(each);
        } catch (e) {
            // Continue
        }
    });

const corpusData: SciFactData[] = fs
    .readFileSync("scifact_corpus.jsonl")
    .toString()
    .split("\n")
    .map((each) => {
        try {
            return JSON.parse(each);
        } catch (e) {
            // Continue
        }
    })
    .slice(0, 10); // Comment this out to use the full corpus

const batchSize = 100;

async function processData(): Promise<void> {
    for (let i = 0; i < corpusData.length; i += batchSize) {
        const batchData = corpusData.slice(i, i + batchSize);
        for (const row of batchData) {
            (await scifactCorpusCollection).add({
                ids: row.doc_id.toString(),
                documents: `${row.title}. ${row.abstract.join(" ")}`,
                metadatas: { structured: row.structured },
            });
        }
    }
}

processData().then(() => null);

const buildPromptWithContext = (claim: string, context: string[]): Array<{
    role: "system" | "user";
    content: string;
}> => [
    {
        role: "system",
        content:
            "I will ask you to assess whether a particular scientific claim, based on evidence provided. " +
            "Output only the text 'True' if the claim is true, 'False' if the claim is false, or 'NEE' if there's " +
            "not enough evidence.",
    },
    {
        role: "user",
        content: `
            The evidence is the following:

            ${context.join(" ")}

            Assess the following claim on the basis of the evidence. Output only the text 'True' if the claim is true,
            'False' if the claim is false, or 'NEE' if there's not enough evidence. Do not output any other text.

            Claim:
            ${claim}

            Assessment:
        `,
    },
];

async function assessClaims(claims: string[]): Promise<string[]> {
    const claimQueryResult = await (
        await scifactCorpusCollection
    ).query({
        queryTexts: claims,
        include: ["documents", "distances"],
        nResults: 3,
    });
    const responses: string[] = [];

    for (let i = 0; i < claimQueryResult.documents.length; i++) {
        const claim = claims[i];
        const context = claimQueryResult.documents[i];
        if (context.length === 0) {
            responses.push("NEE");
            continue;
        }

        const response = await openai_client.chat.completions.create({
            model: "gpt-4o-mini",
            messages: buildPromptWithContext(claim, context),
            max_tokens: 3,
        });

        const formattedResponse = response.choices[0].message.content?.replace(
            "., ",
            "",
        );
        console.log("Claim: ", claim);
        console.log("Response: ", formattedResponse);
        responses.push(formattedResponse ?? "NEE");
    }

    return responses;
}

const tracedAssessClaims = tracer.traceFunction()(assessClaims);

const tracedMain = async (): Promise<void> => {
    const samples = claimData.slice(0, 2); // Get a sample of 2 claims
    await tracedAssessClaims(samples.map((sample) => sample.claim ?? ""));
};
await tracedMain();